from copy import deepcopy
from functools import partial
from unittest.mock import patch

import torch
from torch import nn
from sglang.srt.models import deepseek_ocr as vision
from sglang.srt.models import unlimited_ocr as base
from sglang.srt.models import unlimited_ocr_v2 as dash
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.models.unlimited_ocr_v2 import UnlimitedOCRV2ForCausalLM as UpstreamModel


class ConfigAttention(dash.UnlimitedOCRV2Attention):
    def __init__(self, config, layer_id, quant_config=None, prefix=""):
        # Use the pinned constructor and forward, with a separate residual width.
        attention_config = deepcopy(config)
        attention_config.hidden_size = config.num_attention_heads * config.head_dim
        super().__init__(attention_config, layer_id, quant_config, prefix)
        self.hidden_size = config.hidden_size
        outputs = {
            "q_proj": config.num_attention_heads * self.head_dim,
            "k_proj": config.num_key_value_heads * self.head_dim,
            "v_proj": config.num_key_value_heads * self.head_dim,
            "index_q_proj": self.idx_head_dim,
            "index_k_proj": self.idx_head_dim,
        }
        for name, output_size in outputs.items():
            setattr(self, name, ColumnParallelLinear(config.hidden_size, output_size,
                bias=False, gather_output=name.startswith("index_"), quant_config=quant_config,
                prefix=f"{prefix}.{name}" if prefix else name))
        self.o_proj = RowParallelLinear(config.num_attention_heads * self.head_dim,
            config.hidden_size, bias=False, quant_config=quant_config,
            prefix=f"{prefix}.o_proj" if prefix else "o_proj")


class UnlimitedOCRV2ForCausalLM(UpstreamModel):
    def __init__(self, *, config, quant_config=None, prefix=""):
        specs = config.vision_config.width

        def sam():
            spec = specs["sam_vit_b"]
            encoder = vision.ImageEncoderViT(embed_dim=spec["width"], depth=spec["layers"],
                num_heads=spec["heads"], global_attn_indexes=spec["global_attn_indexes"],
                out_chans=spec["neck_channels"], window_size=14, use_rel_pos=True,
                norm_layer=partial(nn.LayerNorm, eps=1e-6))
            c1, c2 = spec["downsample_channels"]
            encoder.net_2 = nn.Conv2d(spec["neck_channels"], c1, 3, stride=2, padding=1, bias=False)
            encoder.net_3 = nn.Conv2d(c1, c2, 3, stride=2, padding=1, bias=False)
            return encoder

        def clip():
            spec = specs["clip-l-14-224"]
            cfg = deepcopy(vision.vit_model_cfg)
            cfg.update(num_layers=spec["layers"], hidden_size=spec["width"], num_heads=spec["heads"],
                       num_attention_heads=spec["heads"], ffn_hidden_size=spec["width"] * 4)
            return vision.VitModel(cfg)

        with (patch.object(base, "build_sam_vit_b", sam),
              patch.object(base, "build_clip_l", clip),
              patch.object(dash, "UnlimitedOCRV2Attention", ConfigAttention)):
            super().__init__(config=config, quant_config=quant_config, prefix=prefix)


EntryClass = UnlimitedOCRV2ForCausalLM