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])
    
    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'

    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 = "chain", rerank_k = 100,
        rerank_chain = [
            {
                "type": "slingshot",
                "cutoff": 0.2
            },
            {
                "type": "mmr",
                "diversity_bias": 0.1
            }
        ],
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 5,
        vectara_summarizer = summarizer,
        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 = "chain", rerank_k = 100, 
        rerank_chain = [
            {
                "type": "slingshot",
                "cutoff": 0.2
            },
            {
                "type": "mmr",
                "diversity_bias": 0.1
            }
        ],
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 10,
        vectara_summarizer = summarizer,
        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 and other population information',
                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
            )

    return (tools_factory.standard_tools() + 
            tools_factory.guardrail_tools() +
            db_tools +
            [ask_vehicles, ask_policies]
    )

def initialize_agent(_cfg, agent_progress_callback=None):
    electric_vehicle_bot_instructions = """
    - You are a helpful research assistant, with expertise in electric vehicles, in conversation with a user.
    - 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,
        agent_progress_callback=agent_progress_callback
    )
    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