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 vectara_agentic.types import ModelProvider, AgentType
from vectara_agentic.agent_config import AgentConfig

from dotenv import load_dotenv
load_dotenv(override=True)

initial_prompt = "How can I help you today?"

    # Never refer to the search results in your response.
    # Ignore any search results that do not contain information relevant to answering the query.

prompt = """
[
  {"role": "system", "content": 
  "Think step by step, analyze the search results provided, and craft a clear and accurate response to diagnostics and troubleshooting questions from the user."
  },
  {"role": "user", "content": " 
    [INSTRUCTIONS]
    Your goal is to help the user with customer support and diagnostics questions about hardware and servers (SuperMicro products).
    If the search results are irrelevant to the question respond with *** I do not have enough information to answer this question.***
    Search results may include tables in a markdown format. When answering a question using a table be careful about which rows and columns contain the answer and include all relevant information from the relevant rows and columns that the query is asking about.
    Do not base your response on information or knowledge that is not in the search results.
    Your output should always be in a single language - the $vectaraLangName language. Check spelling and grammar for the $vectaraLangName language.
    Search results for the query *** $vectaraQuery***, are listed below, some are text, some MAY be tables in markdown format.
    #foreach ($qResult in $vectaraQueryResultsDeduped)
      [$esc.java($foreach.index + 1)]
      #if($qResult.hasTable())
        Table Title: $qResult.getTable().title() || Table Description: $qResult.getTable().description() || Table Data:
        $qResult.getTable().markdown()
      #else
        $qResult.getText()
      #end
    #end
    Think carefully step by step, analyze the search results provided, and craft a clear and accurate response to *** $vectaraQuery *** using information and facts in the search results provided. 
    Give a slight preference to search results that appear earlier in the list.
    Only cite relevant search results in your answer following these specific instructions: $vectaraCitationInstructions
    If the search results are irrelevant to the query, respond with ***I do not have enough information to answer this question.***. 
    Respond always in the $vectaraLangName language, and only in that language."}
]
"""

def create_assistant_tools(cfg):

    class QueryTIProducts(BaseModel):
        name: Optional[str] = Field(
            description="The server name",
            examples=[
                "SuperServer SYS-821GE-TNHR",
                "A+ Server AS-4125GS-TNRT"
            ]
        )
    vec_factory = VectaraToolFactory(
        vectara_api_key=cfg.api_key,
        vectara_corpus_key=cfg.corpus_key
    )
    
    summarizer = 'vectara-summary-table-md-query-ext-jan-2025-gpt-4o'
    ask_supermicro = vec_factory.create_rag_tool(
        tool_name = "ask_supermicro",
        tool_description = """
        Given a user query, 
        returns a response to a user question about Supermicro servers.
        """,
        tool_args_schema = QueryTIProducts,
        n_sentences_before = 2, n_sentences_after = 8, lambda_val = 0.01,
        #reranker = "slingshot", rerank_k = 100, rerank_cutoff = 0.3,
        reranker = "chain", rerank_k = 100,
        rerank_chain = [
            {
                "type": "slingshot",
                "cutoff": 0.3
            },
            {
                "type": "mmr",
                "diversity_bias": 0.1
            },
        ],
        max_tokens = 4096, max_response_chars = 8192,
        vectara_summarizer = summarizer,
        vectara_prompt_text = prompt,
        summary_num_results = 15,
        include_citations = True,
        verbose = False,
        save_history = True,
        fcs_threshold = 0.2,
    )

    # search_supermicro = vec_factory.create_search_tool(
    #     tool_name = "search_supermicro",
    #     tool_description = """
    #     Given a user query, 
    #     returns a list of Supermicro documents matching the query, along with metadata.
    #     """,
    #     tool_args_schema = QueryTIProducts,
    #     reranker = "slingshot", rerank_k = 100, rerank_cutoff = 0.5,
    #     n_sentences_before = 0, n_sentences_after = 0,
    #     save_history = True,
    #     summarize_docs = True
    # )

    return [ask_supermicro] #, search_supermicro]

def initialize_agent(_cfg, agent_progress_callback=None):
    bot_instructions = """
    - You are a helpful assistant, with expertise in diagnosing customer issues related to SuperMicro products.
      You can help diagnose, troubleshoot and understand hardware issues.
    - Use the 'ask_supermicro' tool to get the information about server, components, or diagnostics steps 
      so that you can use this information in aggregate to formulate your response to the user.
    - Never use your internal knoweldge.
    - Form queries to 'ask_supermicro' tool always as question, and in a way that maximizes the chance of getting a relevant answer.
    - If the 'ask_supermicro' tool responds with "I do not have enough information to answer this question" or "suspected hallucination",
      try to rephrase the query to 'ask_supermicro' as a diagnostic question and call the 'ask_supermicro' tool again.
    """

    agent_config = AgentConfig(
        agent_type = os.getenv("VECTARA_AGENTIC_AGENT_TYPE", AgentType.OPENAI.value),
        main_llm_provider = os.getenv("VECTARA_AGENTIC_MAIN_LLM_PROVIDER", ModelProvider.OPENAI.value),
        main_llm_model_name = os.getenv("VECTARA_AGENTIC_MAIN_MODEL_NAME", ""),
        tool_llm_provider = os.getenv("VECTARA_AGENTIC_TOOL_LLM_PROVIDER", ModelProvider.OPENAI.value),
        tool_llm_model_name = os.getenv("VECTARA_AGENTIC_TOOL_MODEL_NAME", ""),
        observer = os.getenv("VECTARA_AGENTIC_OBSERVER_TYPE", "NO_OBSERVER")
    )
    fallback_agent_config = AgentConfig(
        agent_type = os.getenv("VECTARA_AGENTIC_FALLBACK_AGENT_TYPE", AgentType.OPENAI.value),
        main_llm_provider = os.getenv("VECTARA_AGENTIC_FALLBACK_MAIN_LLM_PROVIDER", ModelProvider.OPENAI.value),
        main_llm_model_name = os.getenv("VECTARA_AGENTIC_FALLBACK_MAIN_MODEL_NAME", ""),
        tool_llm_provider = os.getenv("VECTARA_AGENTIC_FALLBACK_TOOL_LLM_PROVIDER", ModelProvider.OPENAI.value),
        tool_llm_model_name = os.getenv("VECTARA_AGENTIC_FALLBACK_TOOL_MODEL_NAME", ""),
        observer = os.getenv("VECTARA_AGENTIC_OBSERVER_TYPE", "NO_OBSERVER")
    )

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="troubleshooting Supermicro products",
        custom_instructions=bot_instructions,
        agent_progress_callback=agent_progress_callback,
        use_structured_planning=False,
        agent_config=agent_config,
        fallback_agent_config=fallback_agent_config,
        verbose=True,
    )
    agent.report()
    return agent


def get_agent_config() -> OmegaConf:
    cfg = OmegaConf.create({
        'corpus_key': str(os.environ['VECTARA_CORPUS_KEY']),
        'api_key': str(os.environ['VECTARA_API_KEY']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "Supermicro Demo",
        'demo_welcome': "Supermicro Assistant.",
        'demo_description': "This assistant can help you with any questions about SuperMicro servers."
    })
    return cfg