"""Qwen3-VL prompt encoding for Iris-3B, ported from the reference
``iris3b/text/qwen3_vl.py``.

The conditioning is a fixed-length (``text_len`` = 300) stack of 12 Qwen3-VL
hidden layers per token. The prompt is wrapped in a chat template: the system
prefix is fed through the encoder as context but its hidden states are sliced
off; the caption is truncated to ``text_len - len(suffix)`` tokens BEFORE the
assistant-turn suffix is appended, so the suffix can never be the part that is
cut; the rest is right-padded to ``text_len`` and the pad positions are zeroed.

Fixed-length padding is load-bearing, not a convenience: the trunk's joint
attention sees every text position unmasked (only the text adapter masks pad
keys), so the number of pad tokens is part of the distribution the model was
trained on. Always encode to the full ``text_len``.
"""

from typing import List, Tuple

import torch

# Chat template around every prompt (identical to the reference and to Krea 2).
PROMPT_PREFIX = (
    "<|im_start|>system\n"
    "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial "
    "relationships of the objects and background:<|im_end|>\n"
    "<|im_start|>user\n"
)
PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"

# 1-based decoder layers whose hidden states are stacked (hidden_states[i] is
# the output of layer i; index 0 is the embedding output).
DEFAULT_HIDDEN_LAYERS: Tuple[int, ...] = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
DEFAULT_TEXT_LEN = 300


class IrisPromptEncoder:
    """Tokenizes with the Iris template and runs the frozen Qwen3-VL language
    stack. ``__call__`` returns ``(features, mask)`` with features
    ``(B, text_len, num_layers * hidden)`` (layer axis flattened so the toolkit
    treats it as an ordinary (B, L, D) embedding) and ``mask`` ``(B, text_len)``
    bool, True at real tokens."""

    def __init__(
        self,
        tokenizer,
        text_len: int = DEFAULT_TEXT_LEN,
        hidden_layers: Tuple[int, ...] = DEFAULT_HIDDEN_LAYERS,
    ):
        self.tokenizer = tokenizer
        self.text_len = int(text_len)
        self.hidden_layers = tuple(int(i) for i in hidden_layers)
        prefix_ids = tokenizer.encode(PROMPT_PREFIX, add_special_tokens=False)
        suffix_ids = tokenizer.encode(PROMPT_SUFFIX, add_special_tokens=False)
        if not prefix_ids or not suffix_ids:
            raise ValueError("Iris prompt template produced an empty token prefix or suffix")
        self.prefix_ids = torch.tensor(prefix_ids, dtype=torch.long)
        self.suffix_ids = torch.tensor(suffix_ids, dtype=torch.long)
        self.caption_budget = self.text_len - len(suffix_ids)
        if self.caption_budget < 1:
            raise ValueError(f"text_len={self.text_len} leaves no room for a caption")
        pad_id = tokenizer.pad_token_id
        if pad_id is None:
            pad_id = tokenizer.eos_token_id
        if pad_id is None:
            raise ValueError("tokenizer exposes neither pad_token_id nor eos_token_id")
        self.pad_id = int(pad_id)

    def tokenize(self, prompts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
        """``(input_ids, attention_mask)`` of shape ``(B, len(prefix) + text_len)``."""
        caption_ids = self.tokenizer(prompts, add_special_tokens=False)["input_ids"]
        start = len(self.prefix_ids)
        stop = start + self.text_len
        input_ids = torch.full((len(prompts), stop), self.pad_id, dtype=torch.long)
        attention_mask = torch.zeros((len(prompts), stop), dtype=torch.long)
        input_ids[:, :start] = self.prefix_ids
        n_suffix = len(self.suffix_ids)
        for row, ids in enumerate(caption_ids):
            n = min(len(ids), self.caption_budget)
            if n:
                input_ids[row, start : start + n] = torch.tensor(ids[:n], dtype=torch.long)
            input_ids[row, start + n : start + n + n_suffix] = self.suffix_ids
            attention_mask[row, : start + n + n_suffix] = 1
        return input_ids, attention_mask

    @torch.no_grad()
    def __call__(self, text_encoder, prompts: List[str], dtype=None) -> Tuple[torch.Tensor, torch.Tensor]:
        # the language stack alone: Qwen3VLForConditionalGeneration.model is the
        # inner Qwen3VLModel (skips the lm_head logits over 150k vocab)
        decoder = getattr(text_encoder, "model", text_encoder)
        device = text_encoder.device
        input_ids, attention_mask = self.tokenize(prompts)
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        output = decoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
            return_dict=True,
            output_hidden_states=True,
        )
        hidden_states = output.hidden_states
        if hidden_states is None or len(hidden_states) <= self.hidden_layers[-1]:
            raise RuntimeError("Qwen3-VL decoder did not return the requested hidden layers")
        start = len(self.prefix_ids)
        stop = start + self.text_len
        # (B, T, L, D), pads zeroed; with layer offloading the parameters'
        # device is not the compute device, so follow the hidden states
        embeds = torch.stack([hidden_states[layer][:, start:stop] for layer in self.hidden_layers], dim=2)
        mask = attention_mask[:, start:stop].to(embeds.device)
        embeds = embeds * mask[:, :, None, None].to(embeds.dtype)
        if dtype is not None:
            embeds = embeds.to(dtype)
        embeds = embeds.reshape(embeds.shape[0], embeds.shape[1], -1)
        return embeds, mask.bool()