import logging
from dataclasses import dataclass
from typing import Optional, Union

import torch
import torch.nn.functional as F
from transformers.cache_utils import Cache
from transformers.modeling_outputs import ModelOutput
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration

logger = logging.getLogger(__name__)

IGNORE_INDEX = -100


@dataclass
class Qwen3VLCoTEmbeddingOutput(ModelOutput):
    """CoT embedding 模型的前向输出。

    last_hidden_state / attention_mask 用于 last-token pooling 取 embedding；
    lm_loss 为 CoT 段上的 cross entropy，只在传入 labels 时计算；
    logits / past_key_values 只在生成（return_logits=True）时填充。
    """
    lm_loss: Optional[torch.FloatTensor] = None
    num_cot_tokens: Optional[torch.Tensor] = None
    last_hidden_state: Optional[torch.FloatTensor] = None
    attention_mask: Optional[torch.Tensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None


class Qwen3VLForCoTEmbedding(Qwen3VLForConditionalGeneration):
    """在 Qwen3-VL 上同时输出 CoT 生成 logits 与检索 embedding 的模型。

    继承 Qwen3VLForConditionalGeneration 以复用 lm_head、权重 tie 逻辑以及
    generate()，因此同一份权重既能做 CoT 生成（推理阶段 query 先出 CoT）
    也能做 embedding（对 last_hidden_state 做 last-token pooling）。

    query 侧的输入由 collator 拼成 `prefix + cot + <|im_end|>` 的 teacher forcing
    形式，labels 只在 CoT 段有效；doc 侧不传 labels，只取 embedding。
    """
    accepts_loss_kwargs = False

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_logits: bool = False,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> Qwen3VLCoTEmbeddingOutput:
        outputs = self.model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = outputs.last_hidden_state

        lm_loss, num_cot_tokens = None, None
        if labels is not None:
            lm_loss, num_cot_tokens = self.cot_lm_loss(hidden_states, labels)

        # 训练时不需要全序列 logits（显存开销 = B*L*vocab），只有生成时才算
        logits = None
        if return_logits:
            slice_indices = (
                slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0
                else slice(None)
            )
            logits = self.lm_head(hidden_states[:, slice_indices, :])

        return Qwen3VLCoTEmbeddingOutput(
            lm_loss=lm_loss,
            num_cot_tokens=num_cot_tokens,
            last_hidden_state=hidden_states,
            attention_mask=attention_mask,
            logits=logits,
            past_key_values=outputs.past_key_values,
        )

    def cot_lm_loss(self, hidden_states: torch.Tensor, labels: torch.Tensor):
        """只在 labels != -100 的位置上算 next-token cross entropy。

        先按 label mask 把需要预测的 hidden state 挑出来再过 lm_head，
        避免对整条序列展开 (B, L, vocab_size) 的 logits。
        """
        shift_hidden = hidden_states[:, :-1, :]
        shift_labels = labels[:, 1:]

        keep = shift_labels != IGNORE_INDEX
        num_cot_tokens = keep.sum()
        if num_cot_tokens == 0:
            # 整个 batch 都没有 CoT 监督（例如全部被截断），返回挂在计算图上的 0
            logger.warning("No CoT label token in this batch, cot loss is skipped.")
            return hidden_states.sum() * 0.0, num_cot_tokens

        selected_hidden = shift_hidden[keep]                      # (M, H)
        targets = shift_labels[keep].to(torch.long)               # (M,)
        selected_logits = self.lm_head(selected_hidden).float()   # (M, vocab_size)
        return F.cross_entropy(selected_logits, targets), num_cot_tokens

    def prepare_inputs_for_generation(self, *args, **kwargs):
        # generate() 依赖 outputs.logits，因此生成路径上显式打开 logits 计算
        model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
        model_inputs["return_logits"] = True
        return model_inputs


def pooling_last(hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    """取每条序列最后一个有效 token 的 hidden state 作为 embedding。"""
    flipped = attention_mask.flip(dims=[1])
    last_positions = flipped.argmax(dim=1)
    col = attention_mask.shape[1] - last_positions - 1
    row = torch.arange(hidden_state.shape[0], device=hidden_state.device)
    return hidden_state[row, col]