"""Modular vision adaptation interfaces for CORTEX AI.

The module intentionally does not ship vision weights. A real encoder must be
provided by the caller (for example, a Transformers vision model loaded from a
separate approved checkpoint). The adapter can therefore be tested with a
small local encoder while remaining honest about deployment readiness.
"""

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Protocol, Sequence

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class VisionConfig:
    """Configuration for the image-to-language adaptation layer."""

    vision_hidden_size: int
    cortex_hidden_size: int = 7168
    projector_type: str = "linear"
    image_token_count: int | None = None
    status: str = "experimental / in development"


@dataclass
class MultimodalInput:
    """A language prompt plus projected image tokens for a future fusion path."""

    input_ids: Tensor | None
    image_embeddings: Tensor
    attention_mask: Tensor | None = None
    metadata: dict[str, Any] | None = None


class VisionEncoder(Protocol):
    """Minimal callable contract for an image encoder."""

    def __call__(self, images: Any) -> Tensor: ...


class ImageProcessor(Protocol):
    """Minimal preprocessing contract, compatible with HF processors."""

    def __call__(self, images: Any, **kwargs: Any) -> Any: ...


class VisionProjector(nn.Module):
    """Map encoder features into CORTEX hidden space.

    ``projector_type='linear'`` is a single linear map. ``projector_type='mlp'`` is the
    LLaVA-style two-layer gated projection, whose intermediate width equals the target
    hidden size. Only this module is trained in the first experimental stage.
    """

    def __init__(self, config: VisionConfig) -> None:
        super().__init__()
        if config.projector_type == "linear":
            self.projection: nn.Module = nn.Linear(
                config.vision_hidden_size, config.cortex_hidden_size
            )
        elif config.projector_type == "mlp":
            self.projection = nn.Sequential(
                nn.Linear(config.vision_hidden_size, config.cortex_hidden_size),
                nn.GELU(),
                nn.Linear(config.cortex_hidden_size, config.cortex_hidden_size),
            )
        else:
            raise ValueError(
                f"Unknown projector_type={config.projector_type!r}; "
                "expected 'linear' or 'mlp'"
            )
        self.config = config

    def forward(self, features: Tensor) -> Tensor:
        if features.ndim != 3:
            raise ValueError("vision features must have shape [batch, tokens, hidden]")
        return self.projection(features)


class VisionAdapter(nn.Module):
    """Image -> encoder -> projector, without modifying the CORTEX core."""

    def __init__(self, encoder: VisionEncoder, projector: VisionProjector,
                 processor: ImageProcessor | None = None) -> None:
        super().__init__()
        self.encoder = encoder
        self.projector = projector
        self.processor = processor

    def encode_images(self, images: Any) -> Tensor:
        model_inputs = self.processor(images, return_tensors="pt") if self.processor else images
        # Same BatchFeature caveat as native.py: HF processors return a UserDict, which is
        # not a dict subclass, so it must be tested as a Mapping before unpacking by keyword.
        if isinstance(model_inputs, Mapping):
            features = self.encoder(**model_inputs)
        else:
            features = self.encoder(model_inputs)
        if hasattr(features, "last_hidden_state"):
            features = features.last_hidden_state
        if not isinstance(features, Tensor):
            raise TypeError("vision encoder must return a Tensor or an object with last_hidden_state")
        return self.projector(features)

    def prepare_multimodal_input(self, images: Any, input_ids: Tensor | None = None,
                                 attention_mask: Tensor | None = None) -> MultimodalInput:
        projected = self.encode_images(images)
        return MultimodalInput(input_ids, projected, attention_mask,
                               metadata={"vision_status": "experimental / in development"})


class MultimodalFusion(Protocol):
    """Future contract for inserting projected image tokens into the LM context."""

    def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor: ...


class PrefixFusion(nn.Module):
    """Reference fusion strategy: prepend image tokens to text embeddings."""

    def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor:
        if text_embeddings.ndim != 3 or image_embeddings.ndim != 3:
            raise ValueError("text and image embeddings must be rank-3 tensors")
        if text_embeddings.shape[0] != image_embeddings.shape[0]:
            raise ValueError("text and image batches must match")
        return torch.cat((image_embeddings, text_embeddings), dim=1)


def build_vision_adapter(config: VisionConfig, encoder: VisionEncoder,
                         processor: ImageProcessor | None = None) -> VisionAdapter:
    """Build the adaptation layer around an externally supplied encoder."""
    return VisionAdapter(encoder, VisionProjector(config), processor)