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

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

from dotenv import load_dotenv
load_dotenv(override=True)

initial_prompt = "How can I help you today?"

def create_assistant_tools(cfg):

    class QueryCona(BaseModel):
        query: str = Field(description="The user query.")

    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_cona = vec_factory.create_rag_tool(
        tool_name = "ask_cona",
        tool_description = """
        Given a user query, 
        returns a response to a user question about bottling companies.
        """,
        tool_args_schema = QueryCona,
        reranker = "slingshot", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.01,
        vectara_summarizer = summarizer,
        summary_num_results = 20,
        include_citations = True,
        verbose = False
    )
    return [ask_cona] + ToolsFactory().guardrail_tools()

def initialize_agent(_cfg, agent_progress_callback=None):
    bot_instructions = """
    - You are a helpful assistant, with expertise in products from coca cola and other bottling companies.
    - Use the ask_cona tool to answer most questions about any products related to coca cola.
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="Cona services and coca cola",
        custom_instructions=bot_instructions,
        agent_progress_callback=agent_progress_callback,
    )
    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': "Cona Demo",
        'demo_welcome': "Cona Assistant.",
        'demo_description': "This assistant can help you with any questions about Cona Serices."
    })
    return cfg