"""
Video Foundation Model Backbone Adapter.

Wraps Qwen2.5-VL or Cosmos-Reason2 as the visual backbone for Neon VLA.
Extracts visual features from video frames + language instructions.

Why video models over image models:
- Temporal reasoning (motion, physics, object permanence)
- Natural action prediction (video models learn "what happens next")
- Better sim-to-real transfer with Cosmos world model features
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


@dataclass
class BackboneConfig:
    """Configuration for the video backbone."""

    model_id: str = "Qwen/Qwen2.5-VL-7B-Instruct"
    # Supported: "Qwen/Qwen2.5-VL-7B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct",
    #            "Qwen/Qwen2.5-Omni-7B", "Qwen/Qwen2.5-Omni-3B",
    #            "nvidia/Cosmos-Reason2-2B", "nvidia/Cosmos-Reason2-8B"
    torch_dtype: str = "bfloat16"
    device_map: str = "auto"
    attn_implementation: str = "eager"  # "flash_attention_2" if available
    # Video processing
    fps: int = 4  # Frames per second for video input
    max_frames: int = 16  # Max frames to process
    image_size: int = 224  # Resize images to this size
    # Feature extraction
    hidden_size: int = 0  # 0 = auto-detect from model
    freeze_backbone: bool = True  # Freeze backbone during LoRA training
    # Quantization
    load_in_4bit: bool = True  # 4-bit quantization for training
    use_double_quant: bool = True


class VideoBackbone(nn.Module):
    """Video foundation model backbone for extracting visual-language features.

    Supports two backends:
    1. Qwen2.5-VL — Vision-Language model with video understanding
    2. Cosmos-Reason2 — Physical AI reasoning model

    The backbone:
    1. Encodes video frames into visual tokens
    2. Encodes language instructions into text tokens
    3. Fuses them via cross-attention in the VLM
    4. Returns the fused hidden states for action decoding
    """

    def __init__(self, config: BackboneConfig):
        super().__init__()
        self.config = config
        self._model = None
        self._processor = None
        self._tokenizer = None
        self._hidden_size = config.hidden_size
        self._is_cosmos = "cosmos" in config.model_id.lower()

    @property
    def hidden_size(self) -> int:
        """Hidden dimension of the backbone output."""
        if self._hidden_size > 0:
            return self._hidden_size
        # Auto-detect from model config
        if self._model is not None:
            if hasattr(self._model, "config"):
                return getattr(self._model.config, "hidden_size", 2048)
        # Defaults based on known models
        model_id = self.config.model_id.lower()
        if "7b" in model_id:
            return 3584
        elif "3b" in model_id:
            return 2048
        elif "2b" in model_id:
            return 1536
        elif "8b" in model_id:
            return 4096
        return 2048

    def load(self):
        """Load the backbone model and processor. Call this before forward().

        This is separated from __init__ to allow:
        - Lazy loading (don't load until needed)
        - Config serialization without model weights
        """
        if self._model is not None:
            return  # Already loaded

        logger.info(f"Loading video backbone: {self.config.model_id}")

        dtype_map = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }
        torch_dtype = dtype_map.get(self.config.torch_dtype, torch.bfloat16)

        if self._is_cosmos:
            self._load_cosmos(torch_dtype)
        else:
            self._load_qwen_vl(torch_dtype)

        # Freeze backbone if requested (for LoRA training on action heads)
        if self.config.freeze_backbone and self._model is not None:
            for param in self._model.parameters():
                param.requires_grad = False
            logger.info("Backbone frozen (freeze_backbone=True)")

        # Detect hidden size
        if self._hidden_size == 0 and self._model is not None:
            if hasattr(self._model, "config"):
                self._hidden_size = getattr(self._model.config, "hidden_size", 2048)

        logger.info(f"Backbone loaded: hidden_size={self.hidden_size}")

    def _resolve_attn_implementation(self) -> str:
        """Resolve the best available attention implementation.

        Priority: flash_attention_2 > eager
        SDPA is available but can have issues with 4-bit quantized models.

        Returns:
            The attention implementation string to pass to from_pretrained.
        """
        configured = self.config.attn_implementation

        # If user explicitly set something other than "eager", try it
        if configured == "flash_attention_2":
            try:
                import flash_attn  # noqa: F401
                logger.info(f"Using flash_attention_2 (v{flash_attn.__version__})")
                return "flash_attention_2"
            except ImportError:
                logger.warning("flash_attention_2 requested but flash-attn not installed, falling back to eager")
                return "eager"
        elif configured == "sdpa":
            return "sdpa"
        elif configured == "eager":
            return "eager"

        # "auto" or unknown — try FA2 first
        try:
            import flash_attn  # noqa: F401
            logger.info(f"Auto-detected flash_attention_2 (v{flash_attn.__version__})")
            return "flash_attention_2"
        except ImportError:
            logger.info("flash_attention_2 not available, using eager attention")
            return "eager"

    def _load_qwen_vl(self, torch_dtype: torch.dtype):
        """Load Qwen2.5-VL or Qwen2.5-Omni model.

        Handles both standard VL and Omni (audio-native) variants:
        - Qwen2.5-VL: Standard vision-language, uses Qwen2_5_VLForConditionalGeneration
        - Qwen2.5-Omni: Audio-native multimodal, composite model (Thinker + Talker)
          We extract the .thinker sub-model for feature extraction.
        """
        attn_impl = self._resolve_attn_implementation()

        load_kwargs: Dict[str, Any] = {
            "torch_dtype": torch_dtype,
            "device_map": self.config.device_map,
            "trust_remote_code": True,
            "attn_implementation": attn_impl,
        }

        if self.config.load_in_4bit:
            from transformers import BitsAndBytesConfig

            load_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_use_double_quant=self.config.use_double_quant,
            )

        is_omni = "omni" in self.config.model_id.lower()

        if is_omni:
            # Qwen2.5-Omni is a COMPOSITE model (Thinker + Talker + Encoders)
            # It is NOT a CausalLM — must use AutoModel and extract .thinker
            from transformers import AutoModel, AutoProcessor

            logger.info(f"Loading Qwen2.5-Omni (composite model): {self.config.model_id}")
            full_model = AutoModel.from_pretrained(self.config.model_id, **load_kwargs)

            # Extract the thinker (vision-language reasoning component)
            # The talker handles speech generation, encoders handle audio/video input
            if hasattr(full_model, "thinker"):
                self._model = full_model.thinker
                self._is_omni = True  # Track for encode path
                logger.info("Extracted .thinker sub-model from Omni composite")
            else:
                # Fallback: use full model (may work for some Omni variants)
                self._model = full_model
                logger.warning("Omni model has no .thinker attribute, using full model")

            self._processor = AutoProcessor.from_pretrained(
                self.config.model_id, trust_remote_code=True
            )
            logger.info(f"Loaded Qwen2.5-Omni: {self.config.model_id} (attn={attn_impl})")
        else:
            # Standard Qwen2.5-VL
            from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

            self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                self.config.model_id, **load_kwargs
            )
            self._processor = AutoProcessor.from_pretrained(
                self.config.model_id, trust_remote_code=True
            )
            logger.info(f"Loaded Qwen2.5-VL: {self.config.model_id} (attn={attn_impl})")

    def _load_cosmos(self, torch_dtype: torch.dtype):
        """Load Cosmos-Reason2 model.

        Cosmos-Reason2 is built on Qwen3-VL but fine-tuned for physical world reasoning.
        It excels at physics-grounded predictions, world model understanding, and sim2real.
        """
        from transformers import AutoModel, AutoProcessor

        attn_impl = self._resolve_attn_implementation()

        load_kwargs: Dict[str, Any] = {
            "torch_dtype": torch_dtype,
            "device_map": self.config.device_map,
            "trust_remote_code": True,
            "attn_implementation": attn_impl,
        }

        if self.config.load_in_4bit:
            from transformers import BitsAndBytesConfig

            load_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_use_double_quant=self.config.use_double_quant,
            )

        self._model = AutoModel.from_pretrained(self.config.model_id, **load_kwargs)
        self._processor = AutoProcessor.from_pretrained(
            self.config.model_id, trust_remote_code=True
        )
        logger.info(f"Loaded Cosmos-Reason2: {self.config.model_id} (attn={attn_impl})")

    def encode(
        self,
        images: Optional[List[Any]] = None,
        video_frames: Optional[List[Any]] = None,
        text: Optional[str] = None,
        return_dict: bool = True,
    ) -> Dict[str, torch.Tensor]:
        """Encode visual + language inputs into fused features.

        Args:
            images: List of PIL Images (single-frame mode)
            video_frames: List of PIL Images as video frames
            text: Language instruction
            return_dict: If True, return dict; otherwise return tensor

        Returns:
            Dict with:
                - "hidden_states": (batch, seq_len, hidden_size) fused features
                - "pooled": (batch, hidden_size) pooled representation
        """
        if self._model is None:
            self.load()

        if self._is_cosmos:
            return self._encode_cosmos(images, video_frames, text)
        else:
            return self._encode_qwen_vl(images, video_frames, text)

    def _encode_qwen_vl(
        self,
        images: Optional[List] = None,
        video_frames: Optional[List] = None,
        text: Optional[str] = None,
    ) -> Dict[str, torch.Tensor]:
        """Encode using Qwen2.5-VL."""
        # Build the message content
        content = []

        if video_frames:
            # Video mode: pass frames as video
            content.append(
                {
                    "type": "video",
                    "video": video_frames[: self.config.max_frames],
                    "fps": self.config.fps,
                }
            )
        elif images:
            # Image mode
            for img in images[:4]:  # Max 4 images
                content.append({"type": "image", "image": img})

        if text:
            content.append({"type": "text", "text": text})

        messages = [{"role": "user", "content": content}]

        # Process with Qwen2.5-VL processor
        prompt = self._processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self._processor(text=[prompt], images=images, videos=video_frames, return_tensors="pt")

        # Move to device
        device = next(self._model.parameters()).device
        inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

        # Forward pass — get hidden states
        with torch.no_grad():
            outputs = self._model(**inputs, output_hidden_states=True, return_dict=True)

        last_hidden = outputs.hidden_states[-1]  # (batch, seq_len, hidden_size)
        # Pool: use mean of last hidden states
        pooled = last_hidden.mean(dim=1)  # (batch, hidden_size)

        return {
            "hidden_states": last_hidden,
            "pooled": pooled,
        }

    def _encode_cosmos(
        self,
        images: Optional[List] = None,
        video_frames: Optional[List] = None,
        text: Optional[str] = None,
    ) -> Dict[str, torch.Tensor]:
        """Encode using Cosmos-Reason2."""
        # Cosmos uses similar Qwen3-VL architecture
        # Build messages with video/image content
        content = []

        if video_frames:
            content.append({"type": "video", "video": video_frames[: self.config.max_frames]})
        elif images:
            for img in images[:4]:
                content.append({"type": "image", "image": img})

        if text:
            content.append({"type": "text", "text": text})

        messages = [{"role": "user", "content": content}]

        prompt = self._processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self._processor(text=[prompt], images=images, videos=video_frames, return_tensors="pt")

        device = next(self._model.parameters()).device
        inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self._model(**inputs, output_hidden_states=True, return_dict=True)

        last_hidden = outputs.hidden_states[-1]
        pooled = last_hidden.mean(dim=1)

        return {
            "hidden_states": last_hidden,
            "pooled": pooled,
        }

    def forward(
        self,
        images: Optional[List] = None,
        video_frames: Optional[List] = None,
        text: Optional[str] = None,
    ) -> torch.Tensor:
        """Forward pass — returns pooled features."""
        result = self.encode(images=images, video_frames=video_frames, text=text)
        return result["pooled"]
