# /home/user/app/tools/bioportal_tool.py
from langchain_core.tools import BaseTool
from typing import Type, Optional
from pydantic import BaseModel, Field
# Assuming clinical_nlp.umls_bioportal.search_bioportal_term exists and works
from clinical_nlp.umls_bioportal import search_bioportal_term
from services.logger import app_logger
from services.metrics import log_tool_usage

class BioPortalInput(BaseModel):
    term: str = Field(description="The medical term to search for.")
    ontology: Optional[str] = Field(
        default="SNOMEDCT_US", # Default to SNOMEDCT US Edition
        description=(
            "The specific ontology acronym to search within BioPortal (e.g., SNOMEDCT_US, ICD10CM, RXNORM, MESH, LOINC, NCIT). "
            "Defaults to SNOMEDCT_US if not specified."
        )
    )

class BioPortalLookupTool(BaseTool):
    name: str = "bioportal_lookup"
    description: str = (
        "Use this tool to search for medical terms, codes (like ICD-10, SNOMED CT, RxNorm codes), definitions, and concept details "
        "across a wide range of biomedical ontologies via BioPortal. "
        "Useful when you need information from a specific ontology or want to explore different coding systems. "
        "Input is a dictionary with 'term' and optional 'ontology'. If ontology is not specified by the user, you should default to 'SNOMEDCT_US'."
    )
    args_schema: Type[BaseModel] = BioPortalInput

    def _run(self, term: str, ontology: Optional[str] = "SNOMEDCT_US") -> str:
        # Pydantic ensures 'term' is provided. 'ontology' will have its default if not given.
        target_ontology = ontology or "SNOMEDCT_US" # Ensure a default if None is passed
        app_logger.info(f"BioPortal Tool called with term: '{term}', ontology: '{target_ontology}'")
        log_tool_usage(self.name, {"query_term": term, "ontology": target_ontology})

        if not term or not term.strip():
            return "Error from BioPortal lookup: No search term provided."

        try:
            results = search_bioportal_term(term, ontology=target_ontology)
        except Exception as e:
            app_logger.error(f"Exception during BioPortal search for '{term}' in '{target_ontology}': {e}", exc_info=True)
            return f"Error performing BioPortal lookup for '{term}': An unexpected error occurred during the search."

        if isinstance(results, dict) and "error" in results:
            app_logger.warning(f"BioPortal lookup for '{term}' in '{target_ontology}' returned an error: {results['error']}")
            return f"Error from BioPortal lookup for '{term}' (Ontology: {target_ontology}): {results['error']}"
        
        collection = results.get("collection", []) if isinstance(results, dict) else []
        if collection:
            formatted_results = []
            for item in collection[:3]: # Limit to 3 results for brevity
                pref_label = item.get('prefLabel', 'N/A')
                defs = item.get('definition', [])
                definition_str = ("; ".join(d for d in defs if d) if defs else "No definition available.")[:200] + "..." # Truncate long defs
                cui_list = item.get('cui', [])
                cui_str = ", ".join(cui_list) if cui_list else "N/A"
                
                formatted_results.append(
                    f"- Term: {pref_label} (CUIs: {cui_str}). Definition: {definition_str}"
                )
            if formatted_results:
                return f"BioPortal Search Results for '{term}' (Ontology: {target_ontology}):\n" + "\n".join(formatted_results)
            else: # Collection was present but no usable items after formatting
                return f"No specific items found in BioPortal results for '{term}' (Ontology: {target_ontology}), though the query was successful."

        app_logger.warning(f"No results or unexpected format from BioPortal for '{term}' in '{target_ontology}'. Raw: {str(results)[:200]}")
        return f"No results found in BioPortal for term '{term}' in ontology '{target_ontology}'."

    async def _arun(self, term: str, ontology: Optional[str] = "SNOMEDCT_US") -> str:
        # For simplicity, using sync version. Implement true async if search_bioportal_term is async.
        return self._run(term, ontology)