#/usr/bin/env python3

'''
This is the main agent file for the AIAgents4Pharma.
'''

import logging
import hydra
from langgraph_supervisor import create_supervisor
from langchain_openai import ChatOpenAI
from langchain_core.language_models.chat_models import BaseChatModel
from langgraph.checkpoint.memory import MemorySaver
from ...talk2biomodels.agents.t2b_agent import get_app as get_app_t2b
from ...talk2knowledgegraphs.agents.t2kg_agent import get_app as get_app_t2kg
from ..states.state_talk2aiagents4pharma import Talk2AIAgents4Pharma

# Initialize logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_app(uniq_id, llm_model: BaseChatModel):
    '''
    This function returns the langraph app.
    '''
    if hasattr(llm_model, 'model_name'):
        if llm_model.model_name == 'gpt-4o-mini':
            llm_model = ChatOpenAI(model='gpt-4o-mini',
                                temperature=0,
                                model_kwargs={"parallel_tool_calls": False})
    # Load hydra configuration
    logger.log(logging.INFO, "Launching AIAgents4Pharma_Agent with thread_id %s", uniq_id)
    with hydra.initialize(version_base=None, config_path="../configs"):
        cfg = hydra.compose(config_name='config',
                            overrides=['agents/main_agent=default'])
        cfg = cfg.agents.main_agent
    logger.log(logging.INFO, "System_prompt of T2AA4P: %s", cfg.system_prompt)
    with hydra.initialize(version_base=None, config_path="../../talk2biomodels/configs"):
        cfg_t2b = hydra.compose(config_name='config',
                            overrides=['agents/t2b_agent=default'])
        cfg_t2b = cfg_t2b.agents.t2b_agent
    with hydra.initialize(version_base=None, config_path="../../talk2knowledgegraphs/configs"):
        cfg_t2kg = hydra.compose(config_name='config',
                            overrides=['agents/t2kg_agent=default'])
        cfg_t2kg = cfg_t2kg.agents.t2kg_agent
    system_prompt = cfg.system_prompt
    system_prompt += "\n\nHere is the system prompt of T2B agent\n"
    system_prompt += cfg_t2b.state_modifier
    system_prompt += "\n\nHere is the system prompt of T2KG agent\n"
    system_prompt += cfg_t2kg.state_modifier
    # Create supervisor workflow
    workflow = create_supervisor(
        [
            get_app_t2b(uniq_id, llm_model),    # Talk2BioModels
            get_app_t2kg(uniq_id, llm_model)    # Talk2KnowledgeGraphs
        ],
        model=llm_model,
        state_schema=Talk2AIAgents4Pharma,
        # Full history is needed to extract
        # the tool artifacts
        output_mode="full_history",
        add_handoff_back_messages=True,
        prompt=system_prompt
    )

    # Compile and run
    app = workflow.compile(checkpointer=MemorySaver(),
                           name="AIAgents4Pharma_Agent")

    return app