"""MedGemma model wrapper for LangChain compatibility."""
from typing import Any, List, Optional, Iterator
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.callbacks import CallbackManagerForLLMRun


class ChatMedGemma(BaseChatModel):
    """LangChain wrapper for MedGemma multimodal model."""

    model: Any = None
    processor: Any = None
    model_name: str = "google/medgemma-4b-it"
    device: str = "cuda"
    max_new_tokens: int = 2048
    temperature: float = 1.0
    top_p: float = 0.95
    top_k: int = 64

    def __init__(
        self,
        model_name: str = "google/medgemma-4b-it",
        device: str = "cuda",
        load_in_4bit: bool = True,
        max_new_tokens: int = 2048,
        temperature: float = 1.0,
        top_p: float = 0.95,
        top_k: int = 64,
        **kwargs
    ):
        """Initialize MedGemma model.

        Args:
            model_name: HuggingFace model name
            device: Device to load model on (cuda/cpu)
            load_in_4bit: Whether to use 4-bit quantization
            max_new_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            top_p: Top-p sampling parameter
            top_k: Top-k sampling parameter
        """
        super().__init__(**kwargs)
        self.model_name = model_name
        self.device = device
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k

        # Setup quantization
        if load_in_4bit and device == "cuda":
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
            )
        else:
            quantization_config = None

        # Load model and processor
        print(f"Loading MedGemma model: {model_name}...")
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.model = AutoModelForImageTextToText.from_pretrained(
            model_name,
            device_map=device,
            torch_dtype=torch.bfloat16,
            quantization_config=quantization_config,
            trust_remote_code=True,
        ).eval()

        # Enable sampling by default
        self.model.generation_config.do_sample = True
        print(f"✓ MedGemma model loaded successfully")

    def _convert_messages_to_medgemma_format(self, messages: List[BaseMessage]) -> List[dict]:
        """Convert LangChain messages to MedGemma format."""
        converted_messages = []

        for message in messages:
            if isinstance(message, SystemMessage):
                # MedGemma doesn't have system role, prepend to first user message
                converted_messages.append({
                    "role": "system",
                    "content": [{"type": "text", "text": message.content}]
                })
            elif isinstance(message, HumanMessage):
                content = []

                # Handle multimodal content
                if isinstance(message.content, list):
                    for item in message.content:
                        if isinstance(item, dict):
                            if item.get("type") == "image_url":
                                # Extract image path
                                image_url = item.get("image_url", {})
                                if isinstance(image_url, dict):
                                    url = image_url.get("url", "")
                                else:
                                    url = image_url
                                content.append({"type": "image", "url": url})
                            elif item.get("type") == "text":
                                content.append({"type": "text", "text": item.get("text", "")})
                        elif isinstance(item, str):
                            content.append({"type": "text", "text": item})
                elif isinstance(message.content, str):
                    content = [{"type": "text", "text": message.content}]

                converted_messages.append({"role": "user", "content": content})

            elif isinstance(message, AIMessage):
                converted_messages.append({
                    "role": "assistant",
                    "content": [{"type": "text", "text": message.content}]
                })

        return converted_messages

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate a response from MedGemma."""
        # Convert messages to MedGemma format
        medgemma_messages = self._convert_messages_to_medgemma_format(messages)

        # Apply chat template
        inputs = self.processor.apply_chat_template(
            medgemma_messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(device=self.model.device, dtype=torch.bfloat16)

        # Generate response
        with torch.inference_mode():
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                do_sample=True,
                temperature=self.temperature,
                top_p=self.top_p,
                top_k=self.top_k,
                pad_token_id=self.processor.tokenizer.eos_token_id,
            )

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

        # Create ChatGeneration
        message = AIMessage(content=response_text)
        generation = ChatGeneration(message=message)

        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        """Return type of LLM."""
        return "medgemma"

    @property
    def _identifying_params(self) -> dict:
        """Return identifying parameters."""
        return {
            "model_name": self.model_name,
            "device": self.device,
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
        }