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_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 (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(
                tool_name_prefix = "ev",
                content_description = 'Electric Vehicles in the state of Washington',
                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 (ev_load_data, ev_describe_tables and ev_list_tables) to answer analytical queries.
    - IMPORTANT: When using database tools, always "query SELECT * FROM (table_name) LIMIT 25;" first to figure out the format of the columns and
      then call the tool again to try to answer the user's query.
    - Avoid "SELECT *" queries on tables, as they can be slow, instead craft the correct query to get the required information.
    - 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