#!/usr/bin/env python3

"""
This module contains the `GetAnnotationTool` for fetching species annotations 
based on the provided model and species names.
"""
import math
from typing import List, Annotated, Type
import logging
from pydantic import BaseModel, Field
import basico
import pandas as pd
from langgraph.types import Command
from langgraph.prebuilt import InjectedState
from langchain_core.tools.base import BaseTool
from langchain_core.tools.base import InjectedToolCallId
from langchain_core.messages import ToolMessage
from .load_biomodel import ModelData, load_biomodel
from ..api.uniprot import search_uniprot_labels
from ..api.ols import search_ols_labels
from ..api.kegg import fetch_kegg_annotations

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

ols_ontology_abbreviations = {'pato', 'chebi', 'sbo', 'fma', 'pr','go'}

def prepare_content_msg(species_not_found: List[str],
                        species_without_description: List[str]):
    """
    Prepare the content message.
    """
    content = 'Successfully extracted annotations for the species.'
    if species_not_found:
        content += f'''The following species do not exist, and
                        hence their annotations were not extracted:
                        {', '.join(species_not_found)}.'''
    if species_without_description:
        content += f'''The descriptions for the following species
                        were not found:
                        {", ".join(species_without_description)}.'''
    return content

class GetAnnotationInput(BaseModel):
    """
    Input schema for annotation tool.
    """
    sys_bio_model: ModelData = Field(description="model data")
    tool_call_id: Annotated[str, InjectedToolCallId]
    list_species_names: List[str] = Field(
        default=[],
        description='''List of species names to fetch annotations for.
                      If not provided, annotations for all
                      species in the model will be fetched.'''
    )
    state: Annotated[dict, InjectedState]

class GetAnnotationTool(BaseTool):
    """
    Tool for fetching species annotations based on the provided model and species names.
    """
    name: str = "get_annotation"
    description: str = '''A tool to extract annotations for a list of species names
                        based on the provided model. Annotations include
                        the species name, description, database, ID, link,
                        and qualifier. The tool can handle multiple species
                        in a single invoke.'''
    args_schema: Type[BaseModel] = GetAnnotationInput
    return_direct: bool = False

    def _run(self,
             tool_call_id: Annotated[str, InjectedToolCallId],
             state: Annotated[dict, InjectedState],
             list_species_names: List[str] = None,
             sys_bio_model: ModelData = None) -> str:
        """
        Run the tool.
        """
        logger.info("Running the GetAnnotationTool tool for species %s", list_species_names)

        # Prepare the model object
        sbml_file_path = state['sbml_file_path'][-1] if state['sbml_file_path'] else None
        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)

        # Extract all the species names from the model
        df_species = basico.model_info.get_species(model=model_object.copasi_model)

        if df_species is None:
            # for example this may happen with model 20
            raise ValueError("Unable to extract species from the model.")
        # Fetch annotations for the species names
        list_species_names = list_species_names or df_species.index.tolist()

        (annotations_df,
         species_not_found,
         species_without_description) = self._fetch_annotations(list_species_names)

        # Check if annotations are empty
        # If empty, return a message
        if annotations_df.empty:
            logger.warning("The annotations dataframe is empty.")
            return prepare_content_msg(species_not_found, species_without_description)

        # Process annotations
        annotations_df = self._process_annotations(annotations_df)

        # Prepare the simulated data
        dic_annotations_data = {
            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
            'tool_call_id': tool_call_id,
            'data': annotations_df.to_dict()
        }

        # Update the state with the annotations data
        dic_updated_state_for_model = {}
        for key, value in {
            "model_id": [sys_bio_model.biomodel_id],
            "sbml_file_path": [sbml_file_path],
            "dic_annotations_data": [dic_annotations_data]
        }.items():
            if value:
                dic_updated_state_for_model[key] = value

        return Command(
            update=dic_updated_state_for_model | {
                "messages": [
                    ToolMessage(
                        content=prepare_content_msg(species_not_found,
                                                    species_without_description),
                        artifact=True,
                        tool_call_id=tool_call_id
                    )
                ]
            }
        )

    def _fetch_annotations(self, list_species_names: List[str]) -> tuple:
        """
        Fetch annotations for the given species names from the model.
        In this method, we fetch the MIRIAM annotations for the species names.
        If the annotation is not found, we add the species to the list of
        species not found. If the annotation is found, we extract the descriptions
        from the annotation and add them to the data list.

        Args:
            list_species_names (List[str]): List of species names to fetch annotations for.

        Returns:
            tuple: A tuple containing the annotations dataframe, species not found list,
                   and description not found list.
        """
        species_not_found = []
        description_not_found = []
        data = []

        # Loop through the species names
        for species in list_species_names:
            # Get the MIRIAM annotation for the species
            annotation = basico.get_miriam_annotation(name=species)
            # If the annotation is not found, add the species to the list
            if annotation is None:
                species_not_found.append(species)
                continue

            # Extract the descriptions from the annotation
            descriptions = annotation.get("descriptions", [])

            if descriptions == []:
                description_not_found.append(species)
                continue

            # Loop through the descriptions and add them to the data list
            for desc in descriptions:
                data.append({
                    "Species Name": species,
                    "Link": desc["id"],
                    "Qualifier": desc["qualifier"]
                })

        # Create a dataframe from the data list
        annotations_df = pd.DataFrame(data)

        # Return the annotations dataframe and the species not found list
        return annotations_df, species_not_found, description_not_found

    def _process_annotations(self, annotations_df: pd.DataFrame) -> pd.DataFrame:
        """
        Process annotations dataframe to add additional information.
        In this method, we add a new column for the ID, a new column for the database,
        and a new column for the description. We then reorder the columns and process
        the link to format it correctly.

        Args:
            annotations_df (pd.DataFrame): Annotations dataframe to process.

        Returns:
            pd.DataFrame: Processed annotations dataframe
        """
        logger.info("Processing annotations.")
        # Add a new column for the ID
        # Get the ID from the link key
        annotations_df['Id'] = annotations_df['Link'].str.split('/').str[-1]

        # Add a new column for the database
        # Get the database from the link key
        annotations_df['Database'] = annotations_df['Link'].str.split('/').str[-2]

        # Fetch descriptions for the IDs based on the database type
        # by qyerying the respective APIs
        identifiers = annotations_df[['Id', 'Database']].to_dict(orient='records')
        descriptions = self._fetch_descriptions(identifiers)

        # Add a new column for the description
        # Get the description from the descriptions dictionary
        # based on the ID. If the description is not found, use '-'
        annotations_df['Description'] = annotations_df['Id'].apply(lambda x:
                                                                   descriptions.get(x, '-'))
        # annotations_df.index = annotations_df.index + 1

        # Reorder the columns
        annotations_df = annotations_df[
            ["Species Name", "Description", "Database", "Id", "Link", "Qualifier"]
        ]

        # Process the link to format it correctly
        annotations_df["Link"] = annotations_df["Link"].apply(self._process_link)

        # Return the processed annotations dataframe
        return annotations_df

    def _process_link(self, link: str) -> str:
        """
        Process link to format it correctly.
        """
        for ols_ontology_abbreviation in ols_ontology_abbreviations:
            if ols_ontology_abbreviation +'/' in link:
                link = link.replace(f"{ols_ontology_abbreviation}/", "")
        if "kegg.compound" in link:
            link = link.replace("kegg.compound/", "kegg.compound:")
        return link

    def _fetch_descriptions(self, data: List[dict[str, str]]) -> dict[str, str]:
        """
        Fetch protein names or labels based on the database type.
        """
        logger.info("Fetching descriptions for the IDs.")
        results = {}
        grouped_data = {}

        # In the following loop, we create a dictionary with database as the key
        # and a list of identifiers as the value. If either the database or the
        # identifier is NaN, we set it to None.
        for entry in data:
            identifier = entry.get('Id')
            database = entry.get('Database')
            # Check if database is NaN
            if isinstance(database, float):
                if math.isnan(database):
                    database = None
                    results[identifier or "unknown"] = "-"
            else:
                database = database.lower()
                grouped_data.setdefault(database, []).append(identifier)

        # In the following loop, we fetch the descriptions for the identifiers
        # based on the database type.
        # Constants

        for database, identifiers in grouped_data.items():
            if database == 'uniprot':
                results.update(search_uniprot_labels(identifiers))
            elif database in ols_ontology_abbreviations:
                annotations = search_ols_labels([
                        {"Id": id_, "Database": database}
                        for id_ in identifiers
                    ])
                for identifier in identifiers:
                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
            elif database == 'kegg.compound':
                data = [{"Id": identifier, "Database": "kegg.compound"}
                        for identifier in identifiers]
                annotations = fetch_kegg_annotations(data)
                for identifier in identifiers:
                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
            else:
                for identifier in identifiers:
                    results[identifier] = "-"
        return results
