import torch
import torch.nn as nn
from .ideogram4_dit import Ideogram4Fp8QuantBackend  # noqa: F401

LLM_TOKEN_INDICATOR = 3
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)

_DEFAULT_TEXT_ENCODER_CONFIG = {
    "architectures": ["Qwen3VLModel"],
    "dtype": "bfloat16",
    "image_token_id": 151655,
    "model_type": "qwen3_vl",
    "text_config": {
        "attention_bias": False,
        "attention_dropout": 0.0,
        "bos_token_id": 151643,
        "dtype": "bfloat16",
        "eos_token_id": 151645,
        "head_dim": 128,
        "hidden_act": "silu",
        "hidden_size": 4096,
        "initializer_range": 0.02,
        "intermediate_size": 12288,
        "max_position_embeddings": 262144,
        "model_type": "qwen3_vl_text",
        "num_attention_heads": 32,
        "num_hidden_layers": 36,
        "num_key_value_heads": 8,
        "pad_token_id": None,
        "rms_norm_eps": 1e-06,
        "rope_parameters": {
            "mrope_interleaved": True,
            "mrope_section": [24, 20, 20],
            "rope_theta": 5000000,
            "rope_type": "default",
        },
        "use_cache": True,
        "vocab_size": 151936,
    },
    "tie_word_embeddings": False,
    "transformers_version": "5.8.0",
    "video_token_id": 151656,
    "vision_config": {
        "deepstack_visual_indexes": [8, 16, 24],
        "depth": 27,
        "dtype": "bfloat16",
        "hidden_act": "gelu_pytorch_tanh",
        "hidden_size": 1152,
        "in_channels": 3,
        "initializer_range": 0.02,
        "intermediate_size": 4304,
        "model_type": "qwen3_vl_vision",
        "num_heads": 16,
        "num_position_embeddings": 2304,
        "out_hidden_size": 4096,
        "patch_size": 16,
        "spatial_merge_size": 2,
        "temporal_patch_size": 2,
    },
    "vision_end_token_id": 151653,
    "vision_start_token_id": 151652,
    "ideogram_fp8_weight_only": True,
}


class Ideogram4TextEncoder(nn.Module):
    """Qwen3-VL-8B-Instruct wrapper that extracts hidden states from specific layers."""

    def __init__(self, config_path: str = None, **kwargs) -> None:
        super().__init__()
        from transformers import AutoConfig, AutoModel
        if config_path is None:
            config_kwargs = {k: v for k, v in _DEFAULT_TEXT_ENCODER_CONFIG.items() if k != "model_type"}
            config = AutoConfig.for_model("qwen3_vl", **config_kwargs)
        else:
            config = AutoConfig.from_pretrained(config_path, trust_remote_code=True)
        self.model = AutoModel.from_config(config, trust_remote_code=True)
        self.config = config

    def forward(
        self,
        token_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        text_position_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Extract hidden states from specific layers of Qwen3-VL.

        Args:
            token_ids: (B, L) token ids
            attention_mask: (B, L) attention mask
            text_position_ids: (B, L) position ids for text tokens

        Returns:
            (B, L, hidden_size * num_activation_layers) concatenated hidden states
        """
        from transformers.masking_utils import create_causal_mask

        language_model = self.model.language_model

        inputs_embeds = language_model.embed_tokens(token_ids)

        position_ids_4d = text_position_ids[None, ...].expand(4, text_position_ids.shape[0], -1)
        text_position_ids_4d = position_ids_4d[0]
        mrope_position_ids = position_ids_4d[1:]

        causal_mask = create_causal_mask(
            config=language_model.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=None,
            position_ids=text_position_ids_4d,
        )
        position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)

        tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
        captured: dict[int, torch.Tensor] = {}
        hidden_states = inputs_embeds
        for layer_idx, decoder_layer in enumerate(language_model.layers):
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=text_position_ids_4d,
                past_key_values=None,
                position_embeddings=position_embeddings,
            )
            if layer_idx in tap_set:
                captured[layer_idx] = hidden_states

        selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
        stacked = torch.stack(selected, dim=0)
        stacked = torch.permute(stacked, (1, 2, 3, 0))
        batch_size, seq_len, hidden_size = stacked.shape[:3]
        stacked = stacked.reshape(batch_size, seq_len, -1)

        text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
        stacked = stacked * text_mask
        return stacked.to(torch.float32)