"""Qwen3-VL prompt encoding for Mage-Flow.

The reference implementation (mage_flow/models/modules/text_encoder.py) packs
many prompts into one varlen forward; per-sequence the math reduces to a plain
causal forward with positions ``arange(L)`` replicated on every mrope axis and
the leading system-prompt tokens dropped from the output. ai-toolkit encodes
prompts one at a time (embeds are cached per caption), so that per-sequence
form is implemented here directly on the stock HF Qwen3-VL module.

The DiT consumes the final hidden states ``[L - drop_idx, 2560]``; the pooled
"vec" of the reference is unused by the released checkpoints (vec_type null —
the transformer adds a zero vector), so it is not computed.
"""

from __future__ import annotations

import math
from typing import List, Optional

import torch
import torch.nn.functional as F


# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
# is the number of leading (system prompt) tokens dropped from the encoder
# output before it conditions the DiT.
PROMPT_TEMPLATE = {
    "mage-flow": {
        "template": (
            "<|im_start|>system\nDescribe 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{}<|im_end|>\n<|im_start|>assistant\n"
        ),
        "start_idx": 34,
    },
    "mage-flow-edit": {
        "template": (
            "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture,"
            " objects, background), then explain how the user's text instruction should alter or modify the image. "
            "Generate a new image that meets the user's requirements while maintaining consistency with the original "
            "input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
        ),
        "start_idx": 64,
    },
}

# Fixed image placeholder used at edit training time (one per reference image).
EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"


def edit_prompt_body(instruction: str, num_refs: int) -> str:
    """Training-time multi-reference prompt body:
    ``Image 1: <ph>Image 2: <ph>…{instruction}``."""
    prefix = "".join(
        f"Image {j}: {EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1)
    )
    return prefix + instruction


def patch_qwen_vl_patch_embed(model):
    """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
    linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
    falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
    (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
    Returns the number of patch_embed modules patched. (Same patch as the krea2
    extension / Qwen3VLCaptioner.)"""
    patched = 0
    for module in model.modules():
        proj = getattr(module, "proj", None)
        if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
            proj.stride
        ):

            def fast_forward(hidden_states, _proj=proj):
                w = _proj.weight.reshape(_proj.weight.shape[0], -1)
                x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
                return F.linear(x, w, _proj.bias)

            module.forward = fast_forward
            patched += 1
    return patched


def resize_vl_images(
    images: List[torch.Tensor], max_long_edge: int = 384
) -> List["PIL.Image.Image"]:
    """Prepare reference images for the Qwen3-VL conditioning pass.

    Matches the reference's ``_resize_long_edge`` (cap the long edge at
    ``max_long_edge``, preserving aspect ratio, BICUBIC, never upscale) — the
    MLLM only needs a coarse view of the reference; full-resolution detail
    flows through the VAE reference latents. Input tensors are ``(C, H, W)``
    or ``(1, C, H, W)`` in [0, 1]; output is PIL for the Qwen processor.
    """
    from torchvision.transforms.functional import to_pil_image

    out = []
    for img in images:
        if img.dim() == 4:
            img = img[0]
        img = img.float().clamp(0, 1).cpu()
        pil = to_pil_image(img)
        if max_long_edge and max_long_edge > 0:
            w, h = pil.size
            long_edge = max(w, h)
            if long_edge > max_long_edge:
                scale = max_long_edge / long_edge
                new_w = max(1, int(round(w * scale)))
                new_h = max(1, int(round(h * scale)))
                pil = pil.resize((new_w, new_h), resample=3)  # PIL.Image.BICUBIC
        out.append(pil)
    return out


@torch.no_grad()
def encode_mageflow_prompt(
    text_encoder,  # Qwen3VLForConditionalGeneration
    tokenizer,
    prompt: str,
    template_name: str = "mage-flow",
    max_length: int = 2048,
    images: Optional[list] = None,  # list of PIL images (already VL-resized)
    processor=None,  # AutoProcessor, required when images are given
    dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """Encode one prompt (optionally with reference images) to DiT conditioning.

    Returns ``(L - start_idx, 2560)`` final hidden states with the system
    prompt dropped, matching the reference ``TextEncoder.forward`` txt output.
    """
    info = PROMPT_TEMPLATE[template_name]
    template = info["template"]
    drop_idx = int(info["start_idx"])
    device = next(text_encoder.parameters()).device

    if images:
        if processor is None:
            raise ValueError(
                "encoding reference images requires the Qwen3-VL AutoProcessor"
            )
        formatted = template.format(edit_prompt_body(prompt, len(images)))
        vl = processor(text=[formatted], images=list(images), return_tensors="pt")
        input_ids = vl["input_ids"].to(device)
        extra = {}
        for key in ("pixel_values", "image_grid_thw"):
            if vl.get(key) is not None:
                extra[key] = vl[key].to(device)
    else:
        formatted = template.format(prompt)
        input_ids = tokenizer(
            formatted,
            max_length=max_length + drop_idx,
            truncation=True,
            return_tensors="pt",
        ).input_ids.to(device)
        extra = {}

    # The reference packed encoder gives every sequence positions arange(L)
    # replicated on all mrope axes (even for image tokens); a 2D position_ids
    # is expanded to exactly that by the HF text model.
    position_ids = torch.arange(input_ids.shape[1], device=device).unsqueeze(0)

    out = text_encoder.model(
        input_ids=input_ids,
        position_ids=position_ids,
        **extra,
    )
    hidden = out.last_hidden_state[0]  # (L, D)
    return hidden[drop_idx:].to(dtype)