#/usr/bin/env python3

'''

This is the agent file for the Talk2BioModels agent.

'''

import logging

from typing import Annotated

import hydra

from langchain_core.language_models.chat_models import BaseChatModel

from langgraph.checkpoint.memory import MemorySaver

from langgraph.graph import START, StateGraph

from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState

from ..tools.search_models import SearchModelsTool

from ..tools.get_modelinfo import GetModelInfoTool

from ..tools.simulate_model import SimulateModelTool

from ..tools.custom_plotter import CustomPlotterTool

from ..tools.get_annotation import GetAnnotationTool

from ..tools.ask_question import AskQuestionTool

from ..tools.parameter_scan import ParameterScanTool

from ..tools.steady_state import SteadyStateTool

from ..tools.query_article import QueryArticle

from ..states.state_talk2biomodels import Talk2Biomodels

# 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.

    '''

    def agent_t2b_node(state: Annotated[dict, InjectedState]):

        '''

        This function calls the model.

        '''

        logger.log(logging.INFO, "Calling t2b_agent node with thread_id %s", uniq_id)

        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})

        return response

    # Define the tools

    tools = ToolNode([

                    SimulateModelTool(),

                    AskQuestionTool(),

                    CustomPlotterTool(),

                    SearchModelsTool(),

                    GetModelInfoTool(),

                    SteadyStateTool(),

                    ParameterScanTool(),

                    GetAnnotationTool(),

                    QueryArticle()

                ])

    # Load hydra configuration

    logger.log(logging.INFO, "Load Hydra configuration for Talk2BioModels agent.")

    with hydra.initialize(version_base=None, config_path="../configs"):

        cfg = hydra.compose(config_name='config',

                            overrides=['agents/t2b_agent=default'])

        cfg = cfg.agents.t2b_agent

    logger.log(logging.INFO, "state_modifier: %s", cfg.state_modifier)

    # Create the agent

    model = create_react_agent(

                llm_model,

                tools=tools,

                state_schema=Talk2Biomodels,

                prompt=cfg.state_modifier,

                version='v2',

                checkpointer=MemorySaver()

            )

    # Define a new graph

    workflow = StateGraph(Talk2Biomodels)

    # Define the two nodes we will cycle between

    workflow.add_node("agent_t2b", agent_t2b_node)

    # Set the entrypoint as the first node

    # This means that this node is the first one called

    workflow.add_edge(START, "agent_t2b")

    # Initialize memory to persist state between graph runs

    checkpointer = MemorySaver()

    # Finally, we compile it!

    # This compiles it into a LangChain Runnable,

    # meaning you can use it as you would any other runnable.

    # Note that we're (optionally) passing the memory

    # when compiling the graph

    app = workflow.compile(checkpointer=checkpointer,

                           name="T2B_Agent")

    logger.log(logging.INFO,

               "Compiled the graph with thread_id %s and llm_model %s",

               uniq_id,

               llm_model)

    return app