# Copyright (c) ModelScope Contributors. All rights reserved.
import numpy as np
import os
import torch
from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal

from swift.utils import is_deepspeed_enabled, to_device
from ..base import Template
from ..constant import MLLMTemplateType
from ..register import register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Context, Word, findall
from .utils import ChatmlTemplateMeta


@dataclass
class KeyeTemplateMeta(ChatmlTemplateMeta):
    auto_add_bos: bool = False
    stop_words: List[Word] = field(default_factory=lambda: ['<|endoftext|>'])


class KeyeVLTemplate(Template):
    image_token_id = 151655
    video_token_id = 151656
    placeholder_tokens = ['<|image_pad|>', '<|video_pad|>']

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        from keye_vl_utils import fetch_image, fetch_video
        assert media_type in {'image', 'video'}
        if media_type == 'image':
            inputs.images[index] = fetch_image({'image': inputs.images[index]})
            if getattr(self, 'mode', None) == 'lmdeploy':
                return ['<|vision_start|>', [-100], '<|vision_end|>']
            else:
                return ['<|vision_start|><|image_pad|><|vision_end|>']
        else:
            video = inputs.videos[index]
            video, video_kwargs = fetch_video({'video': video})
            if isinstance(video, torch.Tensor):
                video = video.to(torch.uint8)
            inputs.videos[index] = video
            for k, v in video_kwargs.items():
                inputs.mm_processor_kwargs.setdefault(k, []).append(v)
            return ['<|vision_start|><|video_pad|><|vision_end|>']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        processor = self.processor
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        for media_type in ['images', 'videos']:
            mm_data = getattr(inputs, media_type)
            if mm_data:
                if media_type == 'images':
                    media_token = self.image_token_id
                    media_inputs = processor.image_processor(images=mm_data, return_tensors='pt', do_resize=False)
                    media_grid_thw = media_inputs['image_grid_thw']
                else:
                    split_token = self._tokenize('\n')[0]
                    media_inputs = processor(
                        text=['\n'.join(['<|video_pad|>'] * len(mm_data))],
                        videos=mm_data,
                        return_tensors='pt',
                        **inputs.mm_processor_kwargs)
                    splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token)
                    media_grid_thw = media_inputs['video_grid_thw']
                    media_token = self.video_token_id
                idx_list = findall(input_ids, media_token)
                merge_length = processor.image_processor.merge_size**2

                def _get_new_tokens(i):
                    if media_type == 'images':
                        token_len = (media_grid_thw[i].prod() // merge_length)
                        return [media_token] * token_len
                    else:
                        return splited_tokens[i]

                input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
                                                                    _get_new_tokens)
                encoded.update(media_inputs)

        encoded['input_ids'] = input_ids
        encoded['labels'] = labels
        encoded['loss_scale'] = loss_scale
        return encoded

    def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
        if not self.is_training:
            return inputs
        input_ids = inputs['input_ids']
        pixel_values = inputs.get('pixel_values')
        pixel_values_videos = inputs.get('pixel_values_videos')
        image_grid_thw = inputs.get('image_grid_thw')
        video_grid_thw = inputs.get('video_grid_thw')

        base_model = self.get_base_model(model)
        if hasattr(base_model.model, 'embed_tokens'):
            inputs_embeds = base_model.model.embed_tokens(input_ids)
        else:
            inputs_embeds = base_model.model.language_model.embed_tokens(input_ids)

        # Get dtype from visual model, adapting for KeyeVL model structure
        if hasattr(model.visual, 'get_dtype'):
            dtype = model.visual.get_dtype()
        else:
            dtype = model.visual.dtype

        if pixel_values is None and pixel_values_videos is None:  # plain-text
            if is_deepspeed_enabled():
                from PIL import Image
                images = [Image.new('RGB', (32, 32), (0, 0, 0))]
                media_inputs = self.processor.image_processor(images=images, return_tensors='pt')
                device = input_ids.device
                media_inputs = to_device(media_inputs, device)
                pixel_values = media_inputs['pixel_values'].type(dtype)
                # Convert to 5D format for KeyeVL: [num_patches, 3, 14, 14] -> [1, num_patches, 3, 14, 14]
                pixel_values = pixel_values.unsqueeze(0)

                # KeyeVL requires position_ids when pixel_values is 5D
                num_patches = pixel_values.shape[1]
                position_ids = torch.arange(num_patches, device=device)

                # Create dummy grid that works with mlp_AR
                # Assuming merge_size is 2, we need h and w divisible by merge_size
                merge_size = getattr(self.processor.image_processor, 'merge_size', 2)
                grid_size = int(np.sqrt(num_patches))

                # Adjust grid_size to be divisible by merge_size
                if grid_size % merge_size != 0:
                    grid_size = ((grid_size + merge_size - 1) // merge_size) * merge_size

                # For dummy case, use square layout that's compatible with mlp_AR
                dummy_grid_hw = [(1, grid_size, grid_size)]
                sample_indices = torch.zeros(num_patches, dtype=torch.int64, device=device)
                cu_seqlens = torch.tensor([0, num_patches], dtype=torch.int32, device=device)

                vision_outputs = model.visual(
                    pixel_values=pixel_values,
                    image_grid_thw=dummy_grid_hw,
                    position_ids=position_ids,
                    vision_return_embed_list=True,
                    interpolate_pos_encoding=True,
                    sample_indices=sample_indices,
                    cu_seqlens=cu_seqlens,
                    return_pooler_output=False,
                    use_rope=True,
                    window_size=-1,
                )
                image_embeds = vision_outputs.last_hidden_state
                # Process through projector like in normal cases
                image_embeds = model.mlp_AR(image_embeds, dummy_grid_hw)
                # Concatenate all embeddings
                image_embeds = torch.cat(image_embeds, dim=0)
                inputs_embeds += image_embeds.mean() * 0.
        else:
            if pixel_values is not None:
                pixel_values = pixel_values.type(dtype)
                # KeyeVL expects 5D input: (batch_size, sequence_len, channel, height, width)
                # where sequence_len is the total number of patches from all images
                pixel_values = pixel_values.unsqueeze(0)  # [num_patches, 3, 14, 14] -> [1, num_patches, 3, 14, 14]

                if image_grid_thw is not None:
                    image_grid_hws = []
                    for thw in image_grid_thw:
                        if isinstance(thw, torch.Tensor):
                            thw_tuple = tuple(thw.detach().cpu().numpy().tolist())
                        else:
                            thw_tuple = tuple(thw)
                        image_grid_hws.append(thw_tuple)

                    # Prepare position_ids and other parameters for KeyeVL
                    siglip_position_ids = []
                    sample_indices = []
                    cu_seqlens = [0]

                    for idx, thw_tuple in enumerate(image_grid_hws):
                        numel = np.prod(thw_tuple)
                        image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
                        siglip_position_ids.append(image_position_ids)
                        sample_indices.append(torch.full((numel, ), idx, dtype=torch.int64))
                        cu_seqlens.append(cu_seqlens[-1] + numel)

                    siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(pixel_values.device)
                    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values.device)
                    sample_indices = torch.concat(sample_indices, dim=0).to(pixel_values.device)

                    # Call KeyeVL visual model
                    vision_outputs = model.visual(
                        pixel_values=pixel_values,
                        image_grid_thw=image_grid_hws,
                        position_ids=siglip_position_ids,
                        vision_return_embed_list=True,
                        interpolate_pos_encoding=True,
                        sample_indices=sample_indices,
                        cu_seqlens=cu_seqlens,
                        return_pooler_output=False,
                        use_rope=True,
                        window_size=-1,
                    )
                    image_embeds = vision_outputs.last_hidden_state

                    # Process through projector
                    image_embeds = model.mlp_AR(image_embeds, image_grid_thw)
                    # Concatenate all image embeddings
                    image_embeds = torch.cat(image_embeds, dim=0)
                else:
                    # Fallback for case without grid info
                    num_patches = pixel_values.shape[1]
                    position_ids = torch.arange(num_patches, device=pixel_values.device)
                    vision_outputs = model.visual(pixel_values=pixel_values, position_ids=position_ids)
                    image_embeds = vision_outputs.last_hidden_state.reshape(-1,
                                                                            vision_outputs.last_hidden_state.shape[-1])

                image_mask = (input_ids == model.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

            if pixel_values_videos is not None:
                pixel_values_videos = pixel_values_videos.type(dtype)
                # Same processing for videos: convert to 5D format
                pixel_values_videos = pixel_values_videos.unsqueeze(
                    0)  # [num_patches, 3, 14, 14] -> [1, num_patches, 3, 14, 14]

                if video_grid_thw is not None:
                    video_grid_hws = []
                    for thw in video_grid_thw:
                        if isinstance(thw, torch.Tensor):
                            thw_tuple = tuple(thw.detach().cpu().numpy().tolist())
                        else:
                            thw_tuple = tuple(thw)
                        video_grid_hws.append(thw_tuple)

                    siglip_position_ids = []
                    sample_indices = []
                    cu_seqlens = [0]

                    for idx, thw_tuple in enumerate(video_grid_hws):
                        numel = np.prod(thw_tuple)
                        video_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
                        siglip_position_ids.append(video_position_ids)
                        sample_indices.append(torch.full((numel, ), idx, dtype=torch.int64))
                        cu_seqlens.append(cu_seqlens[-1] + numel)

                    siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(pixel_values_videos.device)
                    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values_videos.device)
                    sample_indices = torch.concat(sample_indices, dim=0).to(pixel_values_videos.device)

                    vision_outputs = model.visual(
                        pixel_values=pixel_values_videos,
                        image_grid_thw=video_grid_hws,
                        position_ids=siglip_position_ids,
                        vision_return_embed_list=True,
                        interpolate_pos_encoding=True,
                        sample_indices=sample_indices,
                        cu_seqlens=cu_seqlens,
                        return_pooler_output=False,
                        use_rope=True,
                        window_size=-1,
                    )
                    video_embeds = vision_outputs.last_hidden_state
                    video_embeds = model.mlp_AR(video_embeds, video_grid_thw)
                    video_embeds = torch.cat(video_embeds, dim=0)
                else:
                    # Fallback for case without grid info
                    num_patches = pixel_values_videos.shape[1]
                    position_ids = torch.arange(num_patches, device=pixel_values_videos.device)
                    vision_outputs = model.visual(pixel_values=pixel_values_videos, position_ids=position_ids)
                    video_embeds = vision_outputs.last_hidden_state.reshape(-1,
                                                                            vision_outputs.last_hidden_state.shape[-1])

                video_mask = (input_ids == model.config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds)
                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

        return {'inputs_embeds': inputs_embeds}

    def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        res = super()._data_collator_mm_data(batch)
        second_per_grid_ts = self.gather_list(batch, 'second_per_grid_ts')
        if second_per_grid_ts:
            res['second_per_grid_ts'] = second_per_grid_ts
        return res


# Register the Keye VL template
register_template(KeyeTemplateMeta(MLLMTemplateType.keye_vl, template_cls=KeyeVLTemplate))


class KeyeVL1_5Template(KeyeVLTemplate):

    def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
        return super(KeyeVLTemplate, self)._post_encode(model, inputs)


register_template(
    KeyeTemplateMeta(
        MLLMTemplateType.keye_vl_1_5, template_cls=KeyeVL1_5Template, default_system='You are a helpful assistant.'))
