#!/usr/bin/env python3

"""
Tool for get model information.
"""

import logging
from typing import Type, Optional, Annotated
from dataclasses import dataclass
import basico
from pydantic import BaseModel, Field
from langchain_core.tools import BaseTool
from langchain_core.messages import ToolMessage
from langchain_core.tools.base import InjectedToolCallId
from langgraph.prebuilt import InjectedState
from langgraph.types import Command
from .load_biomodel import ModelData, load_biomodel

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

@dataclass
class RequestedModelInfo:
    """
    Dataclass for storing the requested model information.
    """
    species: bool = Field(description="Get species from the model.", default=False)
    parameters: bool = Field(description="Get parameters from the model.", default=False)
    compartments: bool = Field(description="Get compartments from the model.", default=False)
    units: bool = Field(description="Get units from the model.", default=False)
    description: bool = Field(description="Get description from the model.", default=False)
    name: bool = Field(description="Get name from the model.", default=False)

class GetModelInfoInput(BaseModel):
    """
    Input schema for the GetModelInfo tool.
    """
    requested_model_info: RequestedModelInfo = Field(description="requested model information")
    sys_bio_model: ModelData = Field(description="model data")
    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 GetModelInfoTool(BaseTool):
    """
    This tool ise used extract model information.
    """
    name: str = "get_modelinfo"
    description: str = """A tool for extracting name,
                    description, species, parameters,
                    compartments, and units from a model."""
    args_schema: Type[BaseModel] = GetModelInfoInput

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

        Args:
            requested_model_info (RequestedModelInfo): The requested model information.
            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.

        Returns:
            Command: The updated state of the tool.
        """
        logger.log(logging.INFO,
                   "Calling get_modelinfo tool %s, %s",
                     sys_bio_model,
                   requested_model_info)
        # print (state, 'state')
        sbml_file_path = state['sbml_file_path'][-1] if len(state['sbml_file_path']) > 0 else None
        model_obj = load_biomodel(sys_bio_model,
                                  sbml_file_path=sbml_file_path)
        dic_results = {}
        # Extract species from the model
        if requested_model_info.species:
            df_species = basico.model_info.get_species(model=model_obj.copasi_model)
            if df_species is None:
                raise ValueError("Unable to extract species from the model.")
            # Convert index into a column
            df_species.reset_index(inplace=True)
            dic_results['Species'] = df_species[
                                        ['name',
                                         'compartment',
                                         'type',
                                         'unit',
                                         'initial_concentration',
                                         'display_name']]
            # Convert this into a dictionary
            dic_results['Species'] = dic_results['Species'].to_dict(orient='records')

        # Extract parameters from the model
        if requested_model_info.parameters:
            df_parameters = basico.model_info.get_parameters(model=model_obj.copasi_model)
            if df_parameters is None:
                raise ValueError("Unable to extract parameters from the model.")
            # Convert index into a column
            df_parameters.reset_index(inplace=True)
            dic_results['Parameters'] = df_parameters[
                                        ['name',
                                         'type',
                                         'unit',
                                         'initial_value',
                                         'display_name']]
            # Convert this into a dictionary
            dic_results['Parameters'] = dic_results['Parameters'].to_dict(orient='records')

        # Extract compartments from the model
        if requested_model_info.compartments:
            df_compartments = basico.model_info.get_compartments(model=model_obj.copasi_model)
            dic_results['Compartments'] = df_compartments.index.tolist()
            dic_results['Compartments'] = ','.join(dic_results['Compartments'])

        # Extract description from the model
        if requested_model_info.description:
            dic_results['Description'] = model_obj.description

        # Extract description from the model
        if requested_model_info.name:
            dic_results['Name'] = model_obj.name

        # Extract time unit from the model
        if requested_model_info.units:
            dic_results['Units'] = basico.model_info.get_model_units(model=model_obj.copasi_model)

        # Prepare the dictionary of updated state for the model
        dic_updated_state_for_model = {}
        for key, value in {
                        "model_id": [sys_bio_model.biomodel_id],
                        "sbml_file_path": [sbml_file_path],
                        }.items():
            if value:
                dic_updated_state_for_model[key] = value

        return Command(
            update=dic_updated_state_for_model|{
                    # update the message history
                    "messages": [
                        ToolMessage(
                            content=dic_results,
                            tool_call_id=tool_call_id
                            )
                        ],
                    }
            )