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

from vectara_agentic.agent import Agent
from vectara_agentic.tools import VectaraToolFactory

from dotenv import load_dotenv
load_dotenv(override=True)

initial_prompt = "How can I help you today?"

def create_assistant_tools(cfg):

    class QueryHMC(BaseModel):
        query: str = Field(description="The user query.")
        ticker: Optional[str] = Field(
            default=None,
            description="The company ticker.",
            examples=['GOOG', 'META']
        )
        year: Optional[str] = Field(
            default=None,
            description="The year of the report.",
            examples=[2020, 2023]
        )
        quarter: Optional[int] = Field(
            default=None,
            description="The quarter of the report.",
            examples=[1, 2, 3, 4]
        )
        filing_type: Optional[str] = Field(
            default=None,
            description="The type of filing.",
            examples=['10K', '10Q']
        )

    vec_factory = VectaraToolFactory(
        vectara_api_key=cfg.api_key,
        vectara_customer_id=cfg.customer_id,
        vectara_corpus_id=cfg.corpus_id
    )
    
    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
    #summarizer = 'vectara-summary-ext-24-05-med-omni'
    ask_hmc = vec_factory.create_rag_tool(
        tool_name = "ask_hmc",
        tool_description = """
        Given a user query, 
        returns a response to a user question about fund management companies.
        """,
        tool_args_schema = QueryHMC,
        reranker = "chain", rerank_k = 100, 
        rerank_chain = [
            {
                "type": "slingshot",
                "cutoff": 0.2
            },
            {
                "type": "mmr",
                "diversity_bias": 0.05,
                "limit": 20
            }
        ],
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        vectara_summarizer = summarizer,
        summary_num_results = 10,
        include_citations = True,
    )
    return [ask_hmc]

def initialize_agent(_cfg, agent_progress_callback=None):
    bot_instructions = """
    - You are a helpful assistant, with expertise in management of public company stock portfolios.
    - Use the 'ask_hmc' tool to answer questions about public company performance, risks, and other financial metrics.
      If the tool responds with "I don't have enough information to answer", try rephrasing the question.
    - Use the year, quarter, filing_type and ticker arguments to the 'ask_hmc' tool to get more specific answers.
    - Note that 10Q reports exist for quarters 1, 2, 3 and for the 4th quarter there is a 10K report.
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="Endowment fund management",
        custom_instructions=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_id': str(os.environ['VECTARA_CORPUS_ID']),
        'api_key': str(os.environ['VECTARA_API_KEY']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "Harvard Management Company",
        'demo_welcome': "Harvard Management Company.",
        'demo_description': "AI Assistant.",
    })
    return cfg