# /home/user/app/tools/umls_tool.py
from langchain_core.tools import BaseTool
from typing import Type, Optional # Added Optional
from pydantic import BaseModel, Field
# Assuming clinical_nlp.umls_bioportal.search_umls_term exists and works
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 specific medical term to search for in UMLS.")
    # Example: search_type: Optional[str] = Field(default="exact", description="Type of search, e.g., 'exact', 'approximate'.")

class UMLSLookupTool(BaseTool):
    name: str = "umls_lookup"
    description: str = (
        "Use this tool to find definitions, Concept Unique Identifiers (CUIs), and related concepts for specific medical terms "
        "from the Unified Medical Language System (UMLS). Best for precise terminology lookup. "
        "Input should be a single medical term (e.g., 'myocardial infarction')."
    )
    args_schema: Type[BaseModel] = UMLSInput

    def _run(self, term: str) -> str: # Add other args from UMLSInput if defined
        app_logger.info(f"UMLS Tool called with term: '{term}'")
        log_tool_usage(self.name, {"query_term": term}) # Log term for metrics

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

        try:
            # Pass additional args like search_type to search_umls_term if your function supports them
            results = search_umls_term(term) # Example: search_umls_term(term, search_type=search_type)
        except Exception as e:
            app_logger.error(f"Exception during UMLS search for '{term}': {e}", exc_info=True)
            return f"Error performing UMLS lookup for '{term}': An unexpected error occurred during the search."

        if isinstance(results, dict) and "error" in results:
            app_logger.warning(f"UMLS lookup for '{term}' returned an error: {results['error']}")
            return f"Error from UMLS lookup for '{term}': {results['error']}"
        
        if isinstance(results, dict) and results.get("results"):
            formatted_results = []
            for res_item in results["results"][:3]: # Limit to top 3 results for brevity
                name = res_item.get('name', 'N/A')
                cui = res_item.get('ui', 'N/A') # Assuming 'ui' is the CUI field from your API response
                # You might want to add definition or semantic type if available
                # definition = res_item.get('definition', 'No definition available.')
                formatted_results.append(f"- Term: {name}, CUI: {cui}")
            
            if formatted_results:
                return f"UMLS Search Results for '{term}':\n" + "\n".join(formatted_results)
            else:
                return f"No specific items found in UMLS results for '{term}', though the query was successful."
        
        app_logger.warning(f"No results or unexpected format from UMLS for '{term}'. Raw: {str(results)[:200]}")
        return f"No distinct results found or unexpected response format from UMLS for '{term}'."

    async def _arun(self, term: str) -> str:
        # For simplicity, using sync version. Implement true async if search_umls_term is async.
        # import asyncio
        # return await asyncio.to_thread(self._run, term)
        return self._run(term)