import os
from typing import Optional
from pydantic import Field, BaseModel
from omegaconf import OmegaConf

from llama_index.core.utilities.sql_wrapper import SQLDatabase
from sqlalchemy import create_engine

from dotenv import load_dotenv
load_dotenv(override=True)

from vectara_agent.agent import Agent
from vectara_agent.tools import ToolsFactory, VectaraToolFactory

def create_assistant_tools(cfg):    

    class QueryElectricCars(BaseModel):
        query: str = Field(description="The user query.")

    vec_factory_1 = VectaraToolFactory(vectara_api_key=cfg.api_keys[0],
                                        vectara_customer_id=cfg.customer_id, 
                                        vectara_corpus_id=cfg.corpus_ids[0])
    
    ask_vehicles = vec_factory_1.create_rag_tool(
        tool_name = "ask_vehicles",
        tool_description = """
        Given a user query, 
        returns a response (str) to a user question about electric vehicles based on online resources.
        You can ask this tool any question about electric cars, including the different types of EVs, how they work, the pros and cons of different models, the environmental impact, and more.
        """,
        tool_args_schema = QueryElectricCars,
        reranker = "multilingual_reranker_v1", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 10,
        vectara_summarizer = 'vectara-summary-ext-24-05-sml',
        include_citations = False,
    )

    vec_factory_2 = VectaraToolFactory(vectara_api_key=cfg.api_keys[1],
                                       vectara_customer_id=cfg.customer_id,
                                       vectara_corpus_id=cfg.corpus_ids[1])
    

    class QueryEVLaws(BaseModel):
        query: str = Field(description="The user query")
        state: Optional[str] = Field(default=None,
                                     description="The two digit state code. Optional.",
                                     examples=['CA', 'US', 'WA'])
        type: Optional[str] = Field(default=None,
                                    description="The type of policy. Optional",
                                    examples = ['Laws and Regulations', 'State Incentives', 'Incentives', 'Utility / Private Incentives', 'Programs'])
        
    

    ask_policies = vec_factory_2.create_rag_tool(
        tool_name = "ask_policies",
        tool_description = """
        Given a user query,
        returns a response (str) to a user question about incentives and regulations about electric vehicles in the United States.
        You can ask this tool any question about laws passed by states or the federal government related to electric vehicles.
        """,
        tool_args_schema = QueryEVLaws,
        reranker = "multilingual_reranker_v1", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 10,
        vectara_summarizer = 'vectara-summary-ext-24-05-sml',
        include_citations = False,
    )

    tools_factory = ToolsFactory()

    return (tools_factory.standard_tools() + 
            tools_factory.guardrail_tools() +
            tools_factory.database_tools(
                content_description = 'Electric Vehicles',
                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
            ) +
            [ask_vehicles, ask_policies]
    )

def initialize_agent(_cfg, update_func=None):
    electric_vehicle_bot_instructions = """
    - You are a helpful research assistant, with expertise in electric vehicles, in conversation with a user.
    - Before answering any user query, get sample data from each table in the database, so that you can understand NULL and unique values for each column.
    - For a query with multiple sub-questions, break down the query into the sub-questions, 
      and make separate calls to the ask_vehicles or ask_policies tool to answer each sub-question,
      then combine the answers to provide a complete response.
    - Use the database tools to answer analytical questions.
      Always run SELECT * FROM (table_name) LIMIT 10; to figure out the format of the columns and their values before trying to answer the user's question.
    - When providing links, try to put the name of the website or source of information for the displayed text. Don't just say 'Source'.
    - Never discuss politics, and always respond politely.
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="Electric vehicles in the United States",
        custom_instructions=electric_vehicle_bot_instructions,
        update_func=update_func
    )
    agent.report()
    return agent


def get_agent_config() -> OmegaConf:
    cfg = OmegaConf.create({
        'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
        'corpus_ids': str(os.environ['VECTARA_CORPUS_IDS']).split(','),
        'api_keys': str(os.environ['VECTARA_API_KEYS']).split(','),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_welcome': "Welcome to the EV Assistant demo.",
        'demo_description': "This assistant can help you learn about electric vehicles in the United States, including how they work, the advantages of purchasing them, and reviews on the top choices.",
    })
    return cfg