from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageDraw, ImageFont, ImageOps
from torch.nn import CrossEntropyLoss
from tqdm import tqdm
from transformers import LogitsProcessorList, TextStreamer
from transformers.modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast
from transformers.utils import logging

from .configuration_deepseek_v2 import DeepseekV2Config
from .deepencoder import MlpProjector, _AttrDict, build_clip_l, build_sam_vit_b
from .modeling_deepseekv2 import (
    DeepseekV2ForCausalLM,
    DeepseekV2Model,
    _build_llama_rotary_embedding,
    _prepare_inputs_for_generation,
)

logger = logging.get_logger(__name__)

torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16


@dataclass
class CausalLMOutputWithPastForMTP(MoeCausalLMOutputWithPast):
    last_hidden_state: Optional[torch.Tensor] = None
    # extend the output for MTP and MoE models
    inputs_embeds: Optional[torch.Tensor] = None
    position_ids: Optional[torch.Tensor] = None
    mtp_hidden_states: Optional[list[torch.Tensor]] = None


class NoEOSTextStreamer(TextStreamer):
    def on_finalized_text(self, text: str, stream_end: bool = False):

        eos_text = self.tokenizer.decode([self.tokenizer.eos_token_id], skip_special_tokens=False)
        text = text.replace(eos_text, "\n")
        print(text, flush=True, end="")


# Match Unlimited-OCR generate() defaults.
_DEFAULT_NGRAM_SIZE = 35
_DEFAULT_NGRAM_WINDOW = 1024
_DEFAULT_NGRAM_WHITELIST = frozenset({128821, 128822})  # <td>, </td>


class SlidingWindowNoRepeatNgramProcessor:
    """Ban tokens that would complete an n-gram already seen in the last ``window`` tokens."""

    def __init__(self, ngram_size, window, whitelist_token_ids=None):
        self.ngram_size = ngram_size
        self.window = window
        self.whitelist = set(whitelist_token_ids) if whitelist_token_ids else set()

    def __call__(self, input_ids, scores):
        for batch_idx in range(input_ids.shape[0]):
            sequence = input_ids[batch_idx].tolist()
            if len(sequence) < self.ngram_size:
                continue
            search_start = max(0, len(sequence) - self.window)
            search_end = len(sequence) - self.ngram_size + 1
            if search_end <= search_start:
                continue
            if self.ngram_size > 1:
                current_prefix = tuple(sequence[-(self.ngram_size - 1) :])
            else:
                current_prefix = tuple()
            banned = set()
            for idx in range(search_start, search_end):
                ngram = sequence[idx : idx + self.ngram_size]
                if self.ngram_size == 1 or tuple(ngram[:-1]) == current_prefix:
                    banned.add(ngram[-1])
            banned.difference_update(self.whitelist)
            for token_id in banned:
                scores[batch_idx, token_id] = float("-inf")
        return scores


class DeepseekOCRConfig(DeepseekV2Config):
    model_type = "deepseek_vl_v2"

    def get_text_config(self, *args, **kwargs):
        # Decoder fields are flat on this object. Nested language_config is a
        # dict dump; vLLM rejects that as text_config (no num_attention_heads).
        return self


class DeepseekOCRModel(DeepseekV2Model):
    config_class = DeepseekOCRConfig

    def __init__(self, config: DeepseekV2Config):
        super(DeepseekOCRModel, self).__init__(config)

        self.sam_model = build_sam_vit_b()
        self.vision_model = build_clip_l()
        # self.conv_2 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=2, stride=2)
        n_embed = 1280
        self.projector = MlpProjector(_AttrDict(projector_type="linear", input_dim=2048, n_embed=n_embed))
        embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
        self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
        self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)

        self.rotary_emb = _build_llama_rotary_embedding(config)

    def sam_vision_fwd(self, img_tensor):
        # Processor images stay float32; SAM/CLIP follow the loaded model dtype.
        img_tensor = img_tensor.to(dtype=next(self.sam_model.parameters()).dtype)
        sam_e = self.sam_model(img_tensor)
        vision_e = self.vision_model(img_tensor, sam_e)
        sam_e = sam_e.flatten(2).permute(0, 2, 1)
        vision_e = vision_e[:, 1:]
        concat_e = torch.cat((vision_e, sam_e), dim=-1)
        return self.projector(concat_e)

    def _sam_vision_fwd_local_global(self, patches: torch.Tensor, image_ori: torch.Tensor):
        """Run SAM+CLIP+projector on local crops and global images.

        When spatial sizes match (typical gundam 1024×1024), batch into one forward
        so SAM/CLIP kernels see a larger batch and avoid a second full tower pass.
        """
        n_local = patches.shape[0]
        if patches.shape[-2:] == image_ori.shape[-2:]:
            fused = self.sam_vision_fwd(torch.cat([patches, image_ori], dim=0))
            return fused[:n_local], fused[n_local:]
        return self.sam_vision_fwd(patches), self.sam_vision_fwd(image_ori)

    def compute_inputs_embeds(
        self,
        input_ids: torch.LongTensor = None,
        images: Optional[torch.FloatTensor] = None,
        images_seq_mask: Optional[torch.FloatTensor] = None,
        images_spatial_crop: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ) -> torch.Tensor:
        """Compute ``inputs_embeds`` with vision features injected (no LM forward).

        Extracted from :meth:`forward` so callers with pre-computed hidden states
        (e.g. from vLLM) can get ``inputs_embeds`` for the MTP head without running
        the full DeepseekV2 language model.
        """
        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        inputs_embeds = inputs_embeds.clone()

        B, seq_len, n_dim = inputs_embeds.shape
        if images is not None and (seq_len != 1 or self.training) and torch.sum(images[0][1]).item() != 0:
            patches, image_ori = zip(*images)
            patches = torch.cat(patches, dim=0)
            image_ori = torch.cat(image_ori, dim=0)

            if torch.sum(patches).item() != 0:
                local_features, global_features = self._sam_vision_fwd_local_global(patches, image_ori)
                img_B, hw, n_dim = global_features.shape
                h = w = int(hw**0.5)

                _, hw2, n_dim2 = local_features.shape
                h2 = w2 = int(hw2**0.5)

                global_features = global_features.view(img_B, h, w, n_dim)
                new_line_img_tkn = self.image_newline[None, None, :].expand(img_B, h, 1, n_dim)
                global_features = torch.cat([global_features, new_line_img_tkn], dim=2)

                global_features_list = []
                from_i = 0
                for i, crop_shape in enumerate(images_spatial_crop):
                    width_crop_num, height_crop_num = crop_shape[0], crop_shape[1]

                    local_features_i = local_features.new_empty(0, n_dim2)
                    has_patches = width_crop_num > 1 or height_crop_num > 1
                    if has_patches:
                        to_i = from_i + width_crop_num * height_crop_num
                        local_features_i = (
                            local_features[from_i:to_i]
                            .view(height_crop_num, width_crop_num, h2, w2, n_dim2)
                            .permute(0, 2, 1, 3, 4)
                            .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2)
                        )
                        local_features_i = torch.cat(
                            [
                                local_features_i,
                                self.image_newline[None, None, :].expand(height_crop_num * h2, 1, n_dim2),
                            ],
                            dim=1,
                        )
                        local_features_i = local_features_i.view(-1, n_dim2)
                        from_i = to_i
                    global_features_i = global_features[i].view(-1, n_dim)
                    visual_sep_token = self.view_seperator[None, :]
                    global_features_i = torch.cat([local_features_i, global_features_i, visual_sep_token], dim=0)
                    global_features_list.append(global_features_i)

                global_local_features = torch.cat(global_features_list, dim=0)
            else:
                global_features = self.sam_vision_fwd(image_ori)

                img_B, hw, n_dim = global_features.shape
                h = w = int(hw**0.5)

                global_features = global_features.view(img_B, h, w, n_dim)
                new_line_img_tkn = self.image_newline[None, None, :].expand(img_B, h, 1, n_dim)

                global_features = torch.cat([global_features, new_line_img_tkn], dim=2)
                global_features = global_features.view(img_B, -1, n_dim)
                visual_sep_token = self.view_seperator[None, None, :].expand(img_B, -1, -1)
                global_local_features = torch.cat([global_features, visual_sep_token], dim=1)
                global_local_features = global_local_features.view(-1, n_dim)

            inputs_flat = inputs_embeds.view(B * seq_len, n_dim)
            mask_flat = images_seq_mask.view(B * seq_len)
            images_flat = global_local_features.view(-1, n_dim)
            n_mask = int(mask_flat.sum().item())
            n_visual = images_flat.shape[0]
            if n_mask != n_visual:
                raise RuntimeError(
                    f"Visual token mismatch: encoder produced {n_visual} tokens, "
                    f"mask has {n_mask} positions (diff={n_visual - n_mask}). "
                    f"B={B}, seq_len={seq_len}, "
                    f"images_spatial_crop={images_spatial_crop.tolist()}"
                )
            inputs_flat[mask_flat] = images_flat.to(inputs_flat.dtype)
            inputs_embeds = inputs_flat.view(B, seq_len, n_dim)

        return inputs_embeds

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        images_seq_mask: Optional[torch.FloatTensor] = None,
        images_spatial_crop: Optional[torch.FloatTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:

        inputs_embeds = self.compute_inputs_embeds(
            input_ids=input_ids,
            images=images,
            images_seq_mask=images_seq_mask,
            images_spatial_crop=images_spatial_crop,
            inputs_embeds=inputs_embeds,
        )

        return super(DeepseekOCRModel, self).forward(
            input_ids=None,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_router_logits=output_router_logits,
            return_dict=return_dict,
            cache_position=cache_position,
        )


class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):

    config_class = DeepseekOCRConfig
    model_class = DeepseekOCRModel
    # FastMTP draft tensors ship in the same checkpoint for vLLM. This class is
    # decoder-only — generate() has no speculative path — so ignore them on load.
    _keys_to_ignore_on_load_unexpected = [
        r"(^|\.)mtp_module\.",
        r"(^|\.)mtp_embed_tokens",
    ]

    def __init__(self, config):
        super(DeepseekV2ForCausalLM, self).__init__(config)
        self.model = self.model_class(config)

        self.vocab_size = config.vocab_size

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded
        if getattr(model.config, "mtp_num_heads", 0):
            logger.warning(
                "FastMTP draft weights (mtp_module.*, mtp_embed_tokens.*) were not loaded. "
                "Transformers generate() does not support MTP speculative decoding; "
                "only the 3B MoE decoder runs. Use vLLM with deepseek_ocr_mtp.register() "
                "for FastMTP."
            )
        return loaded

    def get_model(self):
        return self.model

    def set_sam_gradient_checkpointing(self, mode: str = "off") -> None:
        """SAM GC is independent of LLM ``gradient_checkpointing``.

        Prefer FlexAttention (always on when available) for VRAM; use SAM GC only
        if still OOM. ``global`` checkpoints the 4 global-attn blocks only.
        """
        sam = getattr(self.model, "sam_model", None)
        if sam is None:
            return
        if hasattr(sam, "set_gradient_checkpointing"):
            sam.set_gradient_checkpointing(mode)
        elif hasattr(sam, "gradient_checkpointing"):
            sam.gradient_checkpointing = mode not in (None, "off", False)

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        """Enable GC on the LLM decoder only — does not touch SAM (too slow for step time)."""
        super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

    def gradient_checkpointing_disable(self):
        super().gradient_checkpointing_disable()
        self.set_sam_gradient_checkpointing("off")

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        images_seq_mask: Optional[torch.FloatTensor] = None,
        images_spatial_crop: Optional[torch.FloatTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        skip_logits: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPastForMTP]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        output_router_logits = (
            output_router_logits if output_router_logits is not None else self.config.output_router_logits
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # When True, skip the BT x V lm_head matmul so the trainer can run
        # LigerFusedLinearCrossEntropyLoss on ``last_hidden_state`` instead.
        # Default keeps the previous behaviour (always materialise logits).
        if skip_logits is None:
            skip_logits = False

        outputs = self.model(
            input_ids=input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_router_logits=output_router_logits,
            images=images,
            images_seq_mask=images_seq_mask,
            images_spatial_crop=images_spatial_crop,
            cache_position=cache_position,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        loss = None
        if skip_logits:
            # Trainer owns CE via fused linear; do not materialise logits here.
            # Labels (if any) are ignored — the trainer always pops them before
            # calling forward and computes the loss itself.
            logits = None
        else:
            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
            logits = self.lm_head(hidden_states[:, slice_indices, :])

            # Cast logits to fp32 for numerical stability during training.
            # During inference, the logits are cast back to the original dtype (e.g., bfloat16 or float16) to save memory and speed up generation.
            if self.training:
                logits = logits.float()

            if labels is not None:
                # Shift so that tokens < n predict n
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                # Flatten the tokens
                loss_fct = CrossEntropyLoss()
                shift_logits = shift_logits.view(-1, self.config.vocab_size)
                shift_labels = shift_labels.view(-1)
                # Enable model parallelism
                shift_labels = shift_labels.to(shift_logits.device)
                loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPastForMTP(
            loss=loss,
            aux_loss=None,
            logits=logits,
            last_hidden_state=hidden_states,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_logits,
            inputs_embeds=getattr(outputs, "inputs_embeds", None),
            position_ids=getattr(outputs, "position_ids", None),
        )

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return _prepare_inputs_for_generation(self, input_ids, **kwargs)

    def generate(self, *args, **kwargs):
        ngram_size = kwargs.pop("no_repeat_ngram_size", _DEFAULT_NGRAM_SIZE)
        window_size = kwargs.pop("ngram_window", _DEFAULT_NGRAM_WINDOW)
        whitelist_token_ids = kwargs.pop("whitelist_token_ids", _DEFAULT_NGRAM_WHITELIST)
        processors = kwargs.get("logits_processor")
        if processors is None:
            processors = LogitsProcessorList()
        elif not isinstance(processors, (list, tuple)):
            processors = LogitsProcessorList([processors])
        else:
            processors = LogitsProcessorList(processors)
        already = any(isinstance(p, SlidingWindowNoRepeatNgramProcessor) for p in processors)
        if ngram_size and window_size and not already:
            processors.append(SlidingWindowNoRepeatNgramProcessor(ngram_size, window_size, whitelist_token_ids))
        if processors:
            kwargs["logits_processor"] = processors
        return super().generate(*args, **kwargs)

    def disable_torch_init(self):
        """
        Disable the redundant torch default initialization to accelerate model creation.
        """
        import torch

        setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
        setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)