from langchain.tools import BaseTool
from typing import Type
from pydantic import BaseModel, Field
from clinical_nlp.umls_bioportal import search_umls_term
from services.logger import app_logger
from services.metrics import log_tool_usage

class UMLSInput(BaseModel):
    term: str = Field(description="The medical term to search for in UMLS.")

class UMLSLookupTool(BaseTool):
    name: str = "umls_lookup"
    description: str = (
        "Useful for looking up medical terms, concepts, and their definitions or CUIs "
        "(Concept Unique Identifiers) in the Unified Medical Language System (UMLS). "
        "Provide a specific medical term."
    )
    args_schema: Type[BaseModel] = UMLSInput

    def _run(self, term: str) -> str:
        app_logger.info(f"UMLS Tool called with term: {term}")
        log_tool_usage(self.name)
        results = search_umls_term(term)
        if "error" in results:
            return f"Error from UMLS lookup: {results['error']}"
        # Format results for LLM consumption
        if results.get("results"):
            formatted_results = []
            for res in results["results"][:3]: # Limit to 3 results
                 formatted_results.append(f"- Term: {res.get('name', 'N/A')}, CUI: {res.get('ui', 'N/A')}") # Assuming 'ui' for CUI based on UMLS API
            return "UMLS Results:\n" + "\n".join(formatted_results) if formatted_results else "No results found in UMLS."
        return "No results found or unexpected format from UMLS."

    async def _arun(self, term: str) -> str:
        # For simplicity, using sync version for now
        return self._run(term)