"""NVIDIA NV-Reason-CXR tool for expert chest X-ray analysis."""
from typing import Dict, Optional, Tuple, Type, Any
from pathlib import Path
import torch
from PIL import Image
from pydantic import BaseModel, Field
from transformers import AutoProcessor, AutoModelForImageTextToText

from langchain_core.callbacks import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)
from langchain_core.tools import BaseTool


class NVReasonCXRInput(BaseModel):
    """Input schema for the NV-Reason-CXR Tool."""

    image_path: str = Field(
        ...,
        description="Path to the chest X-ray image file (JPG or PNG)",
    )
    query: str = Field(
        default="Find abnormalities and support devices.",
        description="Question or instruction for analyzing the X-ray (e.g., 'Find abnormalities and support devices', 'Provide differential diagnoses', 'Write a structured report')",
    )
    max_new_tokens: int = Field(
        default=2048,
        description="Maximum number of tokens to generate in response"
    )


class NVReasonCXRTool(BaseTool):
    """Tool for expert chest X-ray analysis using NVIDIA's NV-Reason-CXR model.

    This tool uses NVIDIA's specialized NV-Reason-CXR-3B model for detailed chest X-ray
    analysis, including abnormality detection, support device identification, differential
    diagnoses, and structured report generation.
    """

    name: str = "nv_reason_cxr_analysis"
    description: str = (
        "Expert chest X-ray analysis using NVIDIA's specialized NV-Reason-CXR model. "
        "This tool provides detailed medical reasoning and can: "
        "1) Detect abnormalities and support devices in chest X-rays "
        "2) Provide differential diagnoses "
        "3) Generate structured radiology reports "
        "4) Answer specific questions about chest X-ray findings. "
        "Use this for comprehensive chest X-ray interpretation. "
        "Example input: {'image_path': '/path/to/xray.jpg', 'query': 'Find abnormalities and support devices'}"
    )
    args_schema: Type[BaseModel] = NVReasonCXRInput

    model: Any = None
    processor: Any = None
    device: str = "cuda"

    def __init__(
        self,
        model_path: str = "nvidia/NV-Reason-CXR-3B",
        cache_dir: Optional[str] = None,
        load_in_4bit: bool = True,
        device: Optional[str] = "cuda",
    ):
        """Initialize the NV-Reason-CXR Tool."""
        super().__init__()
        self.device = device

        # Load model following NVIDIA's official demo code EXACTLY
        # Requires transformers==4.56.0 (same as NVIDIA's demo)
        try:
            print(f"Loading NV-Reason-CXR model from {model_path}...")
            print(f"Using device: {self.device}")
            print("Using NVIDIA's exact loading pattern with transformers 4.56.0")

            # Match NVIDIA's demo exactly - requires transformers 4.56.0
            # The dtype parameter works correctly in newer transformers versions
            self.model = AutoModelForImageTextToText.from_pretrained(
                pretrained_model_name_or_path=model_path,
                dtype=torch.bfloat16,
            ).eval().to(self.device)

            self.processor = AutoProcessor.from_pretrained(
                model_path,
                use_fast=True,
            )

            print(f"✓ NV-Reason-CXR model loaded successfully on {self.device}")
        except Exception as e:
            print(f"Error loading NV-Reason-CXR model: {e}")
            print(f"Error type: {type(e).__name__}")
            import traceback
            traceback.print_exc()
            raise

    def _run(
        self,
        image_path: str,
        query: str = "Find abnormalities and support devices.",
        max_new_tokens: int = 2048,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Tuple[Dict[str, Any], Dict]:
        """Analyze a chest X-ray image using NV-Reason-CXR.

        Args:
            image_path: Path to the chest X-ray image file
            query: Question or instruction for analysis
            max_new_tokens: Maximum tokens to generate
            run_manager: Optional callback manager

        Returns:
            Tuple[Dict, Dict]: Output dictionary and metadata dictionary
        """
        try:
            # Load image
            image = Image.open(image_path)
            if image.mode != "RGB":
                image = image.convert("RGB")

            # Prepare messages in chat format
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image"},
                        {"type": "text", "text": query}
                    ]
                }
            ]

            # Apply chat template
            prompt = self.processor.apply_chat_template(
                messages,
                add_generation_prompt=True
            )

            # Prepare inputs
            inputs = self.processor(
                text=prompt,
                images=[image],
                return_tensors="pt"
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            # Generate response
            with torch.inference_mode():
                output_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,  # Deterministic for medical analysis
                    pad_token_id=self.processor.tokenizer.eos_token_id,
                )

            # Decode response
            prompt_length = inputs["input_ids"].shape[-1]
            generated_ids = output_ids[0][prompt_length:]
            response = self.processor.decode(
                generated_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True
            )

            output = {
                "analysis": response,
                "query": query,
            }

            metadata = {
                "image_path": image_path,
                "model": "nvidia/NV-Reason-CXR-3B",
                "device": str(self.device),
                "tokens_generated": len(generated_ids),
                "status": "completed",
            }

            return output, metadata

        except Exception as e:
            output = {
                "error": str(e),
                "analysis": None,
            }
            metadata = {
                "image_path": image_path,
                "status": "failed",
                "error_details": str(e),
            }
            return output, metadata

    async def _arun(
        self,
        image_path: str,
        query: str = "Find abnormalities and support devices.",
        max_new_tokens: int = 2048,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> Tuple[Dict[str, Any], Dict]:
        """Asynchronous version of _run."""
        return self._run(image_path, query, max_new_tokens, run_manager)