from typing import Optional
import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.models.qwen3_vl import Qwen3VLModel
from transformers.utils import logging

from .configuration_ops_colqwen3 import OpsColQwen3Config

logger = logging.get_logger(__name__)


class OpsColQwen3PreTrainedModel(PreTrainedModel):
    config_class = OpsColQwen3Config
    base_model_prefix = "ops_colqwen3"
    supports_gradient_checkpointing = True
    _no_split_modules = ["Qwen3VLVisionBlock", "Qwen3DecoderLayer"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_cache_class = True


class OpsColQwen3Model(OpsColQwen3PreTrainedModel):
    _checkpoint_conversion_mapping = {
        r"^language_model": r"qwen3vl.language_model",
        r"^visual": "qwen3vl.visual",
    }

    def __init__(self, config: OpsColQwen3Config):
        super().__init__(config)
        self.config = config

        self.qwen3vl = Qwen3VLModel(config)
        self.dims = config.text_config.hidden_size
        self.custom_text_proj = nn.Linear(config.text_config.hidden_size, self.dims)

        self.mask_non_image_embeddings = config.mask_non_image_embeddings
        self.post_init()

    @classmethod
    def from_pretrained(cls, *args, config: Optional[OpsColQwen3Config] = None, **kwargs):
        key_mapping = kwargs.pop("key_mapping", None)
        if key_mapping is None:
            key_mapping = getattr(cls, "_checkpoint_conversion_mapping", None)
        dims = None
        if 'dims' in kwargs:
            dims = kwargs.pop('dims')
        elif config is not None:
            dims = config.dims

        model = super().from_pretrained(*args, config=config, **kwargs, key_mapping=key_mapping)
        if dims is not None:
            model.dims = dims
        return model

    def forward(self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, image_grid_thw: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
        has_pixel_values = pixel_values is not None

        if has_pixel_values:
            if image_grid_thw is None:
                raise ValueError("`image_grid_thw` must be provided when `pixel_values` is passed.")
            if not torch.is_tensor(image_grid_thw):
                image_grid_thw = torch.as_tensor(image_grid_thw, device=pixel_values.device)

            offsets = image_grid_thw.prod(dim=1)
            unpadded = [pixel_sequence[: int(offset.item())] for pixel_sequence, offset in zip(pixel_values, offsets)]
            pixel_values = torch.cat(unpadded, dim=0) if unpadded else None

        outputs = self.qwen3vl(
            input_ids=input_ids,
            attention_mask=attention_mask,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            use_cache=False,
            output_hidden_states=True,
            return_dict=True,
        )

        last_hidden_states = outputs.last_hidden_state
        proj = self.custom_text_proj(last_hidden_states)

        if self.dims < self.config.text_config.hidden_size:
            proj = proj[..., : self.dims]

        proj = proj / proj.norm(dim=-1, keepdim=True)

        if attention_mask is not None:
            proj = proj * attention_mask.unsqueeze(-1)

        if has_pixel_values and self.mask_non_image_embeddings and input_ids is not None:
            image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
            proj = proj * image_mask

        return proj

    @property
    def patch_size(self) -> int:
        return self.qwen3vl.visual.config.patch_size

    @property
    def spatial_merge_size(self) -> int:
        return self.qwen3vl.visual.config.spatial_merge_size