# /home/user/app/tools/quantum_treatment_optimizer_tool.py
from langchain_core.tools import BaseTool
from typing import Type, List, Dict, Any, Optional
from pydantic import BaseModel, Field

from services.logger import app_logger
from services.metrics import log_tool_usage

try:
    from quantum.optimizer import optimize_treatment
except ImportError:
    app_logger.warning("Actual 'quantum.optimizer.optimize_treatment' not found. Using mock function for QuantumTreatmentOptimizerTool.")
    def optimize_treatment(patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> Dict[str, Any]:
        import random, time # For mock
        time.sleep(random.uniform(0.5,1.0))
        mock_actions = [f"Mock action for {conditions[0] if conditions else 'general health'} considering {patient_data.get('age', 'N/A')} years old."]
        if current_treatments: mock_actions.append(f"Review interaction with {current_treatments[0]}.")
        return {
            "simulated_optimization_id": f"QO-MOCK-{random.randint(1000,9999)}",
            "suggested_actions": mock_actions,
            "primary_focus_condition": conditions[0] if conditions else "Overall Assessment",
            "confidence_level_simulated": random.uniform(0.7, 0.9),
            "summary_notes": "This is a simulated optimization result. Always consult with medical professionals for actual treatment decisions.",
        }

class QuantumOptimizerInput(BaseModel):
    patient_data: Dict[str, Any] = Field(
        description=(
            "A dictionary of relevant patient characteristics. "
            "Examples: {'age': 55, 'gender': 'Male', 'relevant_labs': {'creatinine': 1.2, 'hbA1c': 7.5}, "
            "'allergies': ['penicillin']}. This should be populated from the overall patient context."
        )
    )
    current_treatments: List[str] = Field(
        description="A list of current medications or therapies (e.g., ['Aspirin 81mg', 'Metformin 500mg OD'])."
    )
    conditions: List[str] = Field(
        description="A list of primary diagnosed conditions or symptoms to be addressed (e.g., ['Type 2 Diabetes', 'Hypertension'])."
    )

class QuantumTreatmentOptimizerTool(BaseTool):
    name: str = "quantum_treatment_optimizer"
    description: str = (
        "A specialized (simulated) tool that uses advanced algorithms to suggest optimized or alternative treatment plans "
        "based on provided patient data, current treatments, and diagnosed conditions. "
        "Use this when seeking novel therapeutic strategies, or to optimize complex polypharmacy. "
        "You MUST provide detailed 'patient_data', 'current_treatments', and 'conditions'."
    )
    args_schema: Type[BaseModel] = QuantumOptimizerInput

    def _format_results_for_llm(self, optimization_output: Dict[str, Any]) -> str:
        if not optimization_output or not isinstance(optimization_output, dict):
            return "The optimizer did not return a structured result or the result was empty."
        summary_lines = ["Quantum Treatment Optimizer Suggestions (Simulated):"]
        actions = optimization_output.get("suggested_actions", [])
        if actions:
            summary_lines.append("  Key Suggested Actions/Considerations:")
            for action_str in actions: summary_lines.append(f"    - {action_str}")
        focus = optimization_output.get("primary_focus_condition")
        if focus: summary_lines.append(f"  Primary Focus: Addressing {focus}.")
        confidence = optimization_output.get("confidence_level_simulated")
        if confidence is not None: summary_lines.append(f"  Simulated Confidence Level: {confidence:.0%}")
        notes = optimization_output.get("summary_notes")
        if notes: summary_lines.append(f"  Summary Notes: {notes}")
        sim_id = optimization_output.get("simulated_optimization_id")
        if sim_id: summary_lines.append(f"  (Simulated Optimization ID: {sim_id})")
        if len(summary_lines) == 1:
            return f"The optimizer processed the request but provided no specific actionable suggestions. Raw data: {str(optimization_output)[:300]}"
        return "\n".join(summary_lines)

    def _run(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str], **kwargs: Any) -> str:
        app_logger.info(
            f"Quantum Optimizer Tool called. Patient Data Keys: {list(patient_data.keys())}, "
            f"Treatments: {current_treatments}, Conditions: {conditions}"
        )
        log_tool_usage(self.name, {"conditions_count": len(conditions), "treatments_count": len(current_treatments)})
        if not patient_data or not conditions:
            missing = [item for item, val in [("'patient_data'", patient_data), ("'conditions'", conditions)] if not val]
            return f"Error: Insufficient information. Missing: {', '.join(missing)}. Provide comprehensive details."
        try:
            optimization_output: Dict[str, Any] = optimize_treatment(
                patient_data=patient_data, current_treatments=current_treatments, conditions=conditions
            )
            app_logger.info(f"Quantum optimizer raw output: {str(optimization_output)[:500]}...")
            return self._format_results_for_llm(optimization_output)
        except ImportError as ie:
            app_logger.error(f"ImportError for quantum.optimizer: {ie}", exc_info=True)
            return "Error: The core optimization module is currently unavailable."
        except Exception as e:
            app_logger.error(f"Unexpected error during quantum optimization: {e}", exc_info=True)
            return f"Error in optimization process: {str(e)}. Ensure input data is correct."

    async def _arun(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str], **kwargs: Any) -> str:
        return self._run(patient_data, current_treatments, conditions, **kwargs)