from langchain.tools import BaseTool
from typing import Type, Optional
from pydantic import BaseModel, Field
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", description="The specific ontology to search within BioPortal (e.g., SNOMEDCT, ICD10, RxNorm). Defaults to SNOMEDCT.")

class BioPortalLookupTool(BaseTool):
    name: str = "bioportal_lookup"
    description: str = (
        "Useful for searching medical terms, codes, and definitions across various ontologies "
        "via BioPortal. Specify the term and optionally the ontology (e.g., SNOMEDCT, ICD10CM, RxNorm)."
    )
    args_schema: Type[BaseModel] = BioPortalInput

    def _run(self, term: str, ontology: Optional[str] = "SNOMEDCT") -> str:
        app_logger.info(f"BioPortal Tool called with term: {term}, ontology: {ontology}")
        log_tool_usage(self.name)
        results = search_bioportal_term(term, ontology=ontology or "SNOMEDCT")
        if "error" in results:
            return f"Error from BioPortal lookup: {results['error']}"
        
        # Format results for LLM consumption
        collection = results.get("collection", [])
        if collection:
            formatted_results = []
            for item in collection[:3]: # Limit to 3 results
                defs = item.get('definition', ['N/A'])
                definition_str = "; ".join(defs) if defs else "N/A"
                formatted_results.append(
                    f"- Term: {item.get('prefLabel', 'N/A')}, Definition: {definition_str}, CUIs: {item.get('cui', [])}"
                )
            return f"BioPortal Results (Ontology: {ontology}):\n" + "\n".join(formatted_results)
        return f"No results found in BioPortal for ontology {ontology}."

    async def _arun(self, term: str, ontology: Optional[str] = "SNOMEDCT") -> str:
        return self._run(term, ontology)