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 QueryCFPBComplaints(BaseModel):
        query: str = Field(description="The user query.")

    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_keys,
                                        vectara_customer_id=cfg.customer_id, 
                                        vectara_corpus_id=cfg.corpus_ids)
    
    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'

    ask_complaints = vec_factory.create_rag_tool(
        tool_name = "ask_complaints",
        tool_description = """
        Given a user query, 
        returns a response to a user question about customer complaints about bank services.
        """,
        tool_args_schema = QueryCFPBComplaints,
        reranker = "multilingual_reranker_v1", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 5,
        vectara_summarizer = summarizer,
        include_citations = False,
    )

    tools_factory = ToolsFactory()

    db_tools = tools_factory.database_tools(
                tool_name_prefix = "cfpb",
                content_description = 'Customer complaints about five banks (Bank of America, Wells Fargo, Capital One, Chase, and CITI Bank)',
                sql_database = SQLDatabase(create_engine('sqlite:///cfpb_database.db')),
            )

    return (tools_factory.standard_tools() + 
            tools_factory.guardrail_tools() +
            db_tools +
            [ask_complaints]
    )

def initialize_agent(_cfg, update_func=None):
    cfpb_complaints_bot_instructions = """
    - You are a helpful research assistant, with expertise in complaints from the Consumer Financial Protection Bureau, in conversation with a user.
    - Before answering any user query, use cfpb_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_complaints tool to answer each sub-question,
      then combine the answers to provide a complete response.
    - Use the database tools (cfpb_load_data, cfpb_describe_tables and cfpb_list_tables) to answer analytical queries.
    - 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 cfpb_load_data tool for a query until you have called cfpb_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="Customer complaints from the Consumer Financial Protection Bureau (CFPB)",
        custom_instructions=cfpb_complaints_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']),
        'api_keys': str(os.environ['VECTARA_API_KEYS']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "cfpb-assistant",
        'demo_welcome': "Welcome to the CFPB Customer Complaints demo.",
        'demo_description': "This assistant can help you gain insights into customer complaints to banks recorded by the Consumer Financial Protection Bureau.",
    })
    return cfg