from langchain.tools import BaseTool
from typing import Type, List, Dict, Any
from pydantic import BaseModel, Field
from quantum.optimizer import optimize_treatment
from services.logger import app_logger
from services.metrics import log_tool_usage

class QuantumOptimizerInput(BaseModel):
    patient_data: Dict[str, Any] = Field(description="Dictionary of relevant patient characteristics (e.g., {'age': 55, 'gender': 'male'}).")
    current_treatments: List[str] = Field(description="List of current medications or therapies (e.g., ['Aspirin 81mg', 'Metformin 500mg']).")
    conditions: List[str] = Field(description="List of diagnosed conditions (e.g., ['Type 2 Diabetes', 'Hypertension']).")

class QuantumTreatmentOptimizerTool(BaseTool):
    name: str = "quantum_treatment_optimizer"
    description: str = (
        "A specialized tool that uses quantum-inspired algorithms to suggest optimized treatment plans. "
        "Provide patient data, current treatments, and diagnosed conditions. "
        "Use this when seeking novel therapeutic strategies or to optimize complex polypharmacy."
    )
    args_schema: Type[BaseModel] = QuantumOptimizerInput

    def _run(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> str:
        app_logger.info(f"Quantum Optimizer Tool called with: {patient_data}, {current_treatments}, {conditions}")
        log_tool_usage(self.name)
        try:
            result = optimize_treatment(patient_data, current_treatments, conditions)
            # Format result for LLM
            # Example: "Optimized suggestions: ..., Confidence: ..., Summary: ..."
            # You might want to pretty-print the dict or convert to a string summary
            return f"Quantum Optimizer Results: {result}" 
        except Exception as e:
            app_logger.error(f"Error in QuantumTreatmentOptimizerTool: {e}")
            return f"Error during quantum optimization: {str(e)}"

    async def _arun(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> str:
        # For simplicity, using sync version for now
        return self._run(patient_data, current_treatments, conditions)