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, text

from dotenv import load_dotenv
load_dotenv(override=True)

from vectara_agentic.agent import Agent
from vectara_agentic.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 to a user question about electric vehicles.
        """,
        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-med-omni',
        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'])
        policy_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 to a user question about electric vehicles incentives and regulations, 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-med-omni',
        include_citations = False,
    )

    tools_factory = ToolsFactory()

    db_tools = tools_factory.database_tools(
                tool_name_prefix = "ev",
                content_description = 'Electric Vehicles in the state of Washington',
                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
            )

    def ev_load_sample_data(table_name: str):
        """
        Given a database table name, returns the first 25 rows of the table.
        """
        ev_load_data = [db_tools[i] for i in range(len(db_tools)) if db_tools[i]._metadata.name == 'ev_load_data'][0]
        return ev_load_data(f"SELECT * FROM {table_name} LIMIT 25")


    return ([tools_factory.create_tool(ev_load_sample_data)] + 
            tools_factory.standard_tools() + 
            tools_factory.guardrail_tools() +
            db_tools +
            [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, use ev_describe_tables to understand schema of each table, and use get_sample_data 
      to 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 (ev_load_data, ev_describe_tables and ev_list_tables) to answer analytical queries.
      If you cannot find the information in one of the tables, try using the other tables in the database.
    - IMPORTANT: When using database_tools, always call the ev_load_sample_data tool with the table you want to query
      to understand the table structure, column naming, and values in the table. Never call the ev_load_data tool for a query until you have called ev_load_sample_data.
    - 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_name': "ev-assistant",
        '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 recent trends based on data in the state of Washington.",
    })
    return cfg