#!/usr/bin/env python3

"""

Tool for searching models based on search query.

"""

from typing import Type, Annotated

import logging

from pydantic import BaseModel, Field

import pandas as pd

from basico import biomodels

from langgraph.types import Command

from langchain_core.tools import BaseTool

from langchain_core.messages import ToolMessage

from langchain_core.tools.base import InjectedToolCallId

# Initialize logger

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

class SearchModelsInput(BaseModel):

    """

    Input schema for the search models tool.

    """

    query: str = Field(description="Search models query", default=None)

    num_query: int = Field(description="Top number of models to search",

                           default=10,

                           le=100)

    tool_call_id: Annotated[str, InjectedToolCallId]

# 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 SearchModelsTool(BaseTool):

    """

    Tool for returning the search results based on the search query.

    """

    name: str = "search_models"

    description: str = "Search for only manually curated models in "

    "the BioMmodels database based on keywords."

    args_schema: Type[BaseModel] = SearchModelsInput

    return_direct: bool = False

    def _run(self,

             tool_call_id: Annotated[str, InjectedToolCallId],

             query: str = None,

             num_query: int = 10) -> dict:

        """

        Run the tool.

        Args:

            query (str): The search query.

            num_query (int): The number of models to search.

            tool_call_id (str): The tool call ID.

        Returns:

            dict: The answer to the question in the form of a dictionary.

        """

        logger.log(logging.INFO, "Searching models with the query and number %s, %s",

                   query, num_query)

        # Search for models based on the query

        search_results = biomodels.search_for_model(query, num_results=num_query)

        # Convert the search results to a pandas DataFrame

        df = pd.DataFrame(search_results)

        # Prepare a message to return

        first_n = min(3, len(search_results))

        content = f"Found {len(search_results)} manually curated models"

        content += f" for the query: {query}."

        # Pass the first 3 models to the LLM

        # to avoid hallucinations

        content += f" Here is the summary of the first {first_n} models:"

        for i in range(first_n):

            content += f"\nModel {i+1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"

        # Return the updated state of the tool

        return Command(

                update={

                    # update the message history

                    "messages": [

                        ToolMessage(

                            content=content,

                            tool_call_id=tool_call_id,

                            artifact={'dic_data': df.to_dict(orient='records')}

                            )

                        ],

                    }

            )