import os
from omegaconf import OmegaConf
from vectara_agentic.agent import Agent
from vectara_agentic.agent_config import AgentConfig
from vectara_agentic.types import ModelProvider, AgentType

from dotenv import load_dotenv
load_dotenv(override=True)

initial_prompt = "How can I help you today?"

def initialize_agent(_cfg, agent_progress_callback=None):

    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.from_corpus(
        vectara_corpus_key=_cfg.corpus_key,
        vectara_api_key=_cfg.api_key,
        tool_name="ask_ucsf_ortho",
        data_description="UCSF Orthopedic Website",
        assistant_specialty="UCSF Orthopedic department, helping users with questions about the department.",
        vectara_reranker="multilingual_reranker_v1", vectara_rerank_k=100,
        vectara_lambda_val=0.005,
        vectara_summarizer="vectara-summary-table-md-query-ext-jan-2025-gpt-4o",
        vectara_summary_num_results=20,
        verbose=True,
        agent_progress_callback=agent_progress_callback,
        agent_config=agent_config,
        fallback_agent_config=fallback_agent_config,
    )
    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': "UCSF Ortho Demo",
        'demo_welcome': "",
        'demo_description': "This assistant can help you with any questions about UCSF Orthopedic department."
    })
    return cfg