#!/usr/bin/env python3

"""
Tool for simulating a model.
"""

import logging
from typing import Type, Annotated
from pydantic import BaseModel, Field
from langgraph.types import Command
from langgraph.prebuilt import InjectedState
from langchain_core.tools import BaseTool
from langchain_core.messages import ToolMessage
from langchain_core.tools.base import InjectedToolCallId
from .load_biomodel import ModelData, load_biomodel
from .load_arguments import ArgumentData, add_rec_events
from .utils import get_model_units

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

class SimulateModelInput(BaseModel):
    """
    Input schema for the SimulateModel tool.
    """
    sys_bio_model: ModelData = Field(description="model data",
                                     default=None)
    arg_data: ArgumentData = Field(description=
                                   """time, species, and reocurring data
                                   as well as the simulation name""",
                                   default=None)
    tool_call_id: Annotated[str, InjectedToolCallId]
    state: Annotated[dict, InjectedState]

# Note: It's important that every field has type hints. BaseTool is a
# Pydantic class and not having type hints can lead to unexpected behavior.
class SimulateModelTool(BaseTool):
    """
    Tool for simulating a model.
    """
    name: str = "simulate_model"
    description: str = "A tool to simulate a biomodel"
    args_schema: Type[BaseModel] = SimulateModelInput

    def _run(self,
        tool_call_id: Annotated[str, InjectedToolCallId],
        state: Annotated[dict, InjectedState],
        sys_bio_model: ModelData = None,
        arg_data: ArgumentData = None
    ) -> Command:
        """
        Run the tool.

        Args:
            tool_call_id (str): The tool call ID. This is injected by the system.
            state (dict): The state of the tool.
            sys_bio_model (ModelData): The model data.
            arg_data (ArgumentData): The argument data.

        Returns:
            str: The result of the simulation.
        """
        logger.log(logging.INFO,
                   "Calling simulate_model tool %s, %s",
                   sys_bio_model,
                   arg_data)
        sbml_file_path = state['sbml_file_path'][-1] if len(state['sbml_file_path']) > 0 else None
        model_object = load_biomodel(sys_bio_model,
                                  sbml_file_path=sbml_file_path)
        # Prepare the dictionary of species data
        # that will be passed to the simulate method
        # of the BasicoModel class
        duration = 100.0
        interval = 10
        dic_species_to_be_analyzed_before_experiment = {}
        if arg_data:
            # Prepare the dictionary of species data
            if arg_data.species_to_be_analyzed_before_experiment is not None:
                dic_species_to_be_analyzed_before_experiment = dict(
                    zip(arg_data.species_to_be_analyzed_before_experiment.species_name,
                        arg_data.species_to_be_analyzed_before_experiment.species_concentration))
            # Add reocurring events (if any) to the model
            if arg_data.reocurring_data is not None:
                add_rec_events(model_object, arg_data.reocurring_data)
            # Set the duration and interval
            if arg_data.time_data is not None:
                duration = arg_data.time_data.duration
                interval = arg_data.time_data.interval
        # Update the model parameters
        model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)
        logger.log(logging.INFO,
                   "Following species/parameters updated in the model %s",
                   dic_species_to_be_analyzed_before_experiment)
        # Simulate the model
        df = model_object.simulate(duration=duration, interval=interval)
        logger.log(logging.INFO, "Simulation results ready with shape %s", df.shape)
        dic_simulated_data = {
            'name': arg_data.experiment_name,
            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
            'tool_call_id': tool_call_id,
            'data': df.to_dict()
        }
        # Prepare the dictionary of updated state
        dic_updated_state_for_model = {}
        for key, value in {
            "model_id": [sys_bio_model.biomodel_id],
            "sbml_file_path": [sbml_file_path],
            "dic_simulated_data": [dic_simulated_data],
            }.items():
            if value:
                dic_updated_state_for_model[key] = value
        # Return the updated state of the tool
        return Command(
                update=dic_updated_state_for_model|{
                # update the message history
                "messages": [
                    ToolMessage(
                        content=f"Simulation results of {arg_data.experiment_name}",
                        tool_call_id=tool_call_id,
                        artifact=get_model_units(model_object)
                        )
                    ],
                }
            )