# Copyright (c) ModelScope Contributors. All rights reserved.
import math
import numpy as np
import torch
import transformers
from dataclasses import dataclass, field
from functools import partial
from packaging import version
from torch import nn
from typing import Any, Dict, List, Literal, Optional

from swift.utils import get_env_args, get_packed_seq_params
from ..base import Template
from ..constant import LLMTemplateType, MLLMTemplateType
from ..register import TemplateMeta, register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Context, Prompt, findall
from ..vision_utils import load_audio, load_video_minicpmv_mplug_owl3
from .llama import Llama3TemplateMeta
from .qwen import Qwen2_5TemplateMeta, Qwen3MixedTemplateMeta, QwenTemplateMeta
from .utils import ChatmlTemplateMeta


@dataclass
class MinicpmTemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: ['<s>{{SYSTEM}}'])
    prompt: Prompt = field(default_factory=lambda: ['<用户>{{QUERY}}<AI>'])
    chat_sep: Optional[Prompt] = field(default_factory=list)
    suffix: Prompt = field(default_factory=lambda: ['</s>'])


register_template(MinicpmTemplateMeta(LLMTemplateType.minicpm))


def _remove_idx(arr: List[int], idx_list: List[int]) -> List[int]:
    res = []
    idx_set = set(idx_list)
    for i, x in enumerate(arr):
        if i not in idx_set:
            res.append(x)
    return res


class MiniCPMVTemplate(Template):
    is_v2_5 = False
    use_model = True
    skip_prompt = False
    placeholder_tokens = ['<unk>']

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        if self.mode == 'vllm':
            return ['(<image>./</image>)\n']
        else:
            return [[-100]]

    async def prepare_lmdeploy_turbomind_inputs(self, inputs: Dict[str, Any]) -> None:
        images = inputs.pop('images', None) or []
        if len(images) == 0:
            return
        input_ids = inputs['input_ids']
        idx_list = findall(input_ids, -100)
        idx_list.insert(0, -1)
        new_input_ids = []
        features = []
        for i in range(len(idx_list) - 1):
            new_input_ids += input_ids[idx_list[i] + 1:idx_list[i + 1]]
            context_list = ['<image>', [-100], '</image>']
            feat = [x.squeeze() for x in images[i]['embeddings'].split(1)]
            grid = images[i].get('grid')
            if len(feat) > 1 and grid is not None:
                context_list.append('<slice>')
                for j in range(grid[1]):
                    if j > 0:
                        context_list.append('\n')
                    for _ in range(grid[0]):
                        context_list += ['<image>', [-100], '</image>']
                context_list.append('</slice>\n')
            new_input_ids += self._encode_context_list(context_list)[0]
            features += feat
        new_input_ids += input_ids[idx_list[-1] + 1:]
        inputs['input_ids'] = new_input_ids
        inputs['images'] = features
        await super().prepare_lmdeploy_turbomind_inputs(inputs)

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        images = inputs.images
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        idx_list = findall(input_ids, -100)
        idx = idx_list[0]
        tgt_sizes = None
        slice_mode = getattr(self.config, 'slice_mode', False)
        if slice_mode:
            if self.is_v2_5:
                image_processor = self.processor.image_processor
                image_inputs = image_processor(images, return_tensors='pt').to(self.model_info.torch_dtype)
                placeholder = image_processor.get_slice_image_placeholder(image_inputs.image_sizes[0][0])
                pixel_values = image_inputs['pixel_values']
                tgt_sizes = image_inputs['tgt_sizes']
            else:
                images, placeholder = self.model.get_slice_image_placeholder(images[0], self.processor)
                pixel_values = [[self.model.transform(img) for img in images]]
            placeholder += '\n'
            placeholder_id = self.processor.encode(placeholder, add_special_tokens=False)
            input_ids = (input_ids[:idx] + placeholder_id + input_ids[idx + 1:])
            if labels is not None:
                labels = (labels[:idx] + [-100] * len(placeholder_id) + labels[idx + 1:])
            input_tensor_ids = torch.tensor(input_ids)
            image_start_idx = torch.where(input_tensor_ids == self.processor.im_start_id)[0]
            image_start_idx += 1
            image_end_idx = torch.where(input_tensor_ids == self.processor.im_end_id)[0]
            valid_image_nums = max(len(image_start_idx), len(image_end_idx))
            image_bound = [
                torch.hstack(
                    [image_start_idx[:valid_image_nums].unsqueeze(-1), image_end_idx[:valid_image_nums].unsqueeze(-1)])
            ]
        else:
            placeholder = '<image>' + '<unk>' * self.config.query_num + '</image>\n'
            placeholder_id = self.processor.encode(placeholder, add_special_tokens=False)
            input_ids = (input_ids[:idx] + placeholder_id + input_ids[idx + 1:])
            if labels is not None:
                labels = (labels[:idx] + [-100] * len(placeholder_id) + labels[idx + 1:])
            image_bound = [torch.tensor([[idx, idx + self.config.query_num]])]
            pixel_values = [[self.model.transform(images[0])]]
        encoded = {
            'input_ids': input_ids,
            'labels': labels,
            'image_bound': image_bound,
            'pixel_values': pixel_values,
            'tgt_sizes': tgt_sizes
        }
        return encoded

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        inputs_embeds, _ = model.get_vllm_embedding(inputs)
        return {'inputs_embeds': inputs_embeds}

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = {}
        for k in ['pixel_values', 'image_bound', 'tgt_sizes']:
            res[k] = self.gather_list(batch, k)
        res.update(super()._data_collator(batch, padding_to=padding_to))
        return res


register_template(MinicpmTemplateMeta(MLLMTemplateType.minicpmv, template_cls=MiniCPMVTemplate))


class MiniCPMV2_5Template(MiniCPMVTemplate):
    is_v2_5 = True


register_template(Llama3TemplateMeta(
    MLLMTemplateType.minicpmv2_5,
    template_cls=MiniCPMV2_5Template,
))


class MiniCPMV2_6Template(MiniCPMVTemplate):

    def init_env_args(self):
        super().init_env_args()
        self.max_num_frames = get_env_args('max_num_frames', int, 64)
        self.max_slice_nums = get_env_args('max_slice_nums', int, None)
        self.video_max_slice_nums = get_env_args('video_max_slice_nums', int, 1)  # or 2

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type in {'image', 'video'}
        load_video = partial(load_video_minicpmv_mplug_owl3, max_num_frames=self.max_num_frames)
        image_context = super().replace_tag('image', index, inputs)
        if media_type == 'image':
            return image_context
        elif media_type == 'video':
            return self.replace_video2image(load_video, inputs, lambda i: image_context)

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = Template._encode(self, inputs)
        images = inputs.images
        use_video = bool(inputs.videos)
        use_image_id = True
        max_slice_nums = self.max_slice_nums
        if use_video:
            max_slice_nums = self.video_max_slice_nums
            use_image_id = False
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        idx_list = findall(input_ids, -100)

        image_processor = self.processor.image_processor
        image_inputs = image_processor([images], return_tensors='pt',
                                       max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype)

        def _get_new_tokens(i):
            placeholder = image_processor.get_slice_image_placeholder(
                image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id)
            placeholder += '\n'
            return self.processor.encode(placeholder, add_special_tokens=False)

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

        if inputs.images:
            input_tensor_ids = torch.tensor(input_ids)
            unk_token = self.processor.encode('<unk>', add_special_tokens=False)[0]
            indices = (input_tensor_ids == unk_token).nonzero(as_tuple=True)[0].tolist()

            ranges = []
            start = indices[0]
            for i in range(1, len(indices)):
                if indices[i] != indices[i - 1] + 1:
                    ranges.append([start, indices[i - 1] + 1])
                    start = indices[i]
            ranges.append([start, indices[-1] + 1])
            image_bound = [torch.tensor(ranges)]
        else:
            image_bound = [[]]

        encoded = {
            'input_ids': input_ids,
            'labels': labels,
            'loss_scale': loss_scale,
            'image_bound': image_bound,
            'pixel_values': image_inputs['pixel_values'],
            'tgt_sizes': image_inputs['tgt_sizes']
        }
        return encoded


register_template(QwenTemplateMeta(
    MLLMTemplateType.minicpmv2_6,
    template_cls=MiniCPMV2_6Template,
))

register_template(ChatmlTemplateMeta(
    MLLMTemplateType.minicpmv4,
    template_cls=MiniCPMV2_6Template,
))

register_template(Qwen2_5TemplateMeta(
    MLLMTemplateType.minicpmo,
    template_cls=MiniCPMV2_6Template,
))


class MiniCPMV4_5Template(MiniCPMV2_6Template):

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = Template._encode(self, inputs)
        images = inputs.images
        use_video = bool(inputs.videos)
        use_image_id = True
        max_slice_nums = self.max_slice_nums
        if use_video:
            max_slice_nums = self.video_max_slice_nums
            use_image_id = False
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        idx_list = findall(input_ids, -100)

        image_processor = self.processor.image_processor
        image_inputs = image_processor([images], return_tensors='pt',
                                       max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype)

        def _get_new_tokens(i):
            placeholder = image_processor.get_slice_image_placeholder(
                image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id)
            placeholder += '\n'
            return self.processor.encode(placeholder, add_special_tokens=False)

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

        if inputs.images:
            input_tensor_ids = torch.tensor(input_ids)
            unk_token = self.processor.encode('<unk>', add_special_tokens=False)[0]
            indices = (input_tensor_ids == unk_token).nonzero(as_tuple=True)[0].tolist()

            ranges = []
            start = indices[0]
            for i in range(1, len(indices)):
                if indices[i] != indices[i - 1] + 1:
                    ranges.append([start, indices[i - 1] + 1])
                    start = indices[i]
            ranges.append([start, indices[-1] + 1])
            image_bound = [torch.tensor(ranges)]
        else:
            image_bound = [[]]

        encoded = {
            'input_ids': input_ids,
            'labels': labels,
            'loss_scale': loss_scale,
            'image_bound': image_bound,
            'pixel_values': image_inputs['pixel_values'],
            'tgt_sizes': image_inputs['tgt_sizes'],
            'temporal_ids': image_inputs['temporal_ids'],
        }
        return encoded

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = {}
        for k in ['pixel_values', 'image_bound', 'tgt_sizes', 'temporal_ids']:
            res[k] = self.gather_list(batch, k)
        res.update(Template._data_collator(self, batch, padding_to=padding_to))
        return res


register_template(
    Qwen3MixedTemplateMeta(
        MLLMTemplateType.minicpmv4_5,
        template_cls=MiniCPMV4_5Template,
        is_thinking=True,
        thinking_prefix='<think>\n',
    ))


class MiniCPMO4_5Template(MiniCPMV4_5Template):
    SAMPLING_RATE = 16000
    MAX_AUDIO_DURATION = 30  # seconds

    def init_env_args(self):
        super().init_env_args()
        self.use_audio_in_video = get_env_args('use_audio_in_video', bool, False)

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        if media_type == 'image' or media_type == 'video' and not self.use_audio_in_video:
            return super().replace_tag(media_type, index, inputs)
        elif media_type == 'audio':
            # Load audio from file path to numpy array at 16kHz
            if isinstance(inputs.audios[index], str):
                inputs.audios[index] = load_audio(inputs.audios[index], sampling_rate=self.SAMPLING_RATE)
            return ['<|audio_start|><|audio_end|>']
        elif media_type == 'video':
            from minicpmo.utils import get_video_frame_audio_segments
            video = inputs.videos[inputs.video_idx]
            video_segments, audio_segments, _ = get_video_frame_audio_segments(
                video, use_audio=self.use_audio_in_video, stack_frames=1)
            # Insert frames into images list at current position
            images = inputs.images
            inputs.images = images[:inputs.image_idx] + video_segments + images[inputs.image_idx:]
            # Build context list
            image_context = [[-100]]
            context_list = []
            if self.use_audio_in_video and audio_segments:
                # Insert audio segments into audios list at current position
                audios = inputs.audios
                inputs.audios = audios[:inputs.audio_idx] + audio_segments + audios[inputs.audio_idx:]
                audio_context = ['<|audio_start|><|audio_end|>']
                # Interleave: one image placeholder + one audio placeholder per second
                for i in range(len(video_segments)):
                    context_list += image_context
                    if i < len(audio_segments):
                        context_list += audio_context
                inputs.audio_idx += len(audio_segments)
            else:
                for _ in range(len(video_segments)):
                    context_list += image_context
            inputs.image_idx += len(video_segments)
            return context_list

    def _get_audio_num_tokens(self, audio_sample_len: int) -> int:
        """Compute the number of <unk> placeholder tokens for an audio of given sample count.

        This mirrors the official get_audio_placeholder logic:
        1. mel frames = ceil(audio_samples / hop_length)
        2. after CNN downsampling: (mel_frames - 1) // 2 + 1
        3. after avg pooling: (cnn_frames - pool_step) // pool_step + 1
        """
        hop_length = self.processor.audio_processor.hop_length  # 160
        pool_step = self.config.audio_pool_step  # 5
        feature_lens = math.ceil(audio_sample_len / hop_length)
        feature_lens_after_cnn = (feature_lens - 1) // 2 + 1
        output_lens = (feature_lens_after_cnn - pool_step) // pool_step + 1
        return output_lens

    def _extract_audio_features(self, audios: List[np.ndarray]):
        """Extract mel features from audio arrays using the WhisperFeatureExtractor.

        Handles chunking of long audios (>30s) into segments.
        Matches the official audio_feature_extract output format.

        Returns:
            audio_features: tensor [N, 80, max_frames] or [] if no audios
            audio_feature_lens: [tensor([l1, l2, ...])] or None
        """
        audio_processor = self.processor.audio_processor
        max_audio_inp_len = self.MAX_AUDIO_DURATION * self.SAMPLING_RATE
        all_audio_features = []
        all_audio_lens = []

        for audio in audios:
            # Chunk long audios at 30s boundaries
            if len(audio) <= max_audio_inp_len:
                chunks = [audio]
            else:
                chunks = [audio[i:i + max_audio_inp_len] for i in range(0, len(audio), max_audio_inp_len)]

            for chunk in chunks:
                audio_input = audio_processor(
                    chunk,
                    sampling_rate=self.SAMPLING_RATE,
                    return_tensors='pt',
                    padding='max_length',
                    return_attention_mask=True,
                )
                feat = audio_input['input_features']  # [1, 80, frames]
                actual_len = audio_input['attention_mask'].sum(dim=1)  # [1]
                feat = feat[:, :, :actual_len[0]]
                all_audio_features.append(feat.squeeze(0))  # [80, actual_frames]
                all_audio_lens.append(actual_len[0])

        if all_audio_features:
            # Pad and stack: [N, 80, max_frames] — same as official processor
            audio_features = torch.nn.utils.rnn.pad_sequence(
                [f.transpose(0, 1) for f in all_audio_features],
                batch_first=True,
                padding_value=0.0,
            ).transpose(1, 2)
            audio_feature_lens = [torch.hstack(all_audio_lens)]
        else:
            audio_features = []
            audio_feature_lens = None

        return audio_features, audio_feature_lens

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        # Step 1: Base encode — produces input_ids with -100 for images
        # and audio_start_id,audio_end_id pairs for audios
        encoded = Template._encode(self, inputs)
        images = inputs.images
        use_video = bool(inputs.videos)
        use_image_id = True
        max_slice_nums = self.max_slice_nums
        if use_video:
            max_slice_nums = self.video_max_slice_nums
            use_image_id = False
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)

        # Step 2: Process images — replace -100 tokens with image placeholders
        idx_list = findall(input_ids, -100)
        image_processor = self.processor.image_processor
        image_inputs = image_processor([images], return_tensors='pt',
                                       max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype)

        def _get_new_tokens(i):
            placeholder = image_processor.get_slice_image_placeholder(
                image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id)
            placeholder += '\n'
            return self.processor.encode(placeholder, add_special_tokens=False)

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

        # Step 3: Process audios — expand audio_start/audio_end pairs with <unk> tokens
        tokenizer = self.processor.tokenizer
        audio_start_id = tokenizer.convert_tokens_to_ids('<|audio_start|>')
        audio_end_id = tokenizer.convert_tokens_to_ids('<|audio_end|>')
        unk_token_id = tokenizer.convert_tokens_to_ids('<unk>')

        audio_features = None
        audio_feature_lens = None

        if inputs.audios:
            audios = inputs.audios
            audio_features, audio_feature_lens = self._extract_audio_features(audios)

            # Find empty audio placeholder pairs (audio_start_id immediately followed by audio_end_id)
            audio_placeholder_positions = []
            for i in range(len(input_ids) - 1):
                if input_ids[i] == audio_start_id and input_ids[i + 1] == audio_end_id:
                    audio_placeholder_positions.append(i)

            assert len(audio_placeholder_positions) == len(audios), \
                f'Found {len(audio_placeholder_positions)} audio placeholders but have {len(audios)} audios'

            # Expand each audio placeholder with <unk> tokens
            offset = 0
            for i, audio in enumerate(audios):
                num_tokens = self._get_audio_num_tokens(len(audio))
                unk_tokens = [unk_token_id] * num_tokens
                pos = audio_placeholder_positions[i] + offset
                # Current: [..., audio_start_id, audio_end_id, ...]
                # Target:  [..., audio_start_id, unk*N, audio_end_id, ...]
                input_ids = input_ids[:pos + 1] + unk_tokens + input_ids[pos + 1:]
                if labels is not None:
                    labels = labels[:pos + 1] + [-100] * num_tokens + labels[pos + 1:]
                if loss_scale is not None:
                    scale_val = loss_scale[pos]
                    loss_scale = loss_scale[:pos + 1] + [scale_val] * num_tokens + loss_scale[pos + 1:]
                offset += num_tokens

        # Step 4: Compute image_bound using start/end token boundaries
        # This is more robust than finding consecutive <unk> tokens, especially
        # when both image and audio use <unk> as placeholder.
        input_tensor_ids = torch.tensor(input_ids)

        if images:
            im_start_id = tokenizer.convert_tokens_to_ids('<image>')
            im_end_id = tokenizer.convert_tokens_to_ids('</image>')
            slice_start_id = tokenizer.convert_tokens_to_ids('<slice>')
            slice_end_id = tokenizer.convert_tokens_to_ids('</slice>')

            start_cond = (input_tensor_ids == im_start_id) | (input_tensor_ids == slice_start_id)
            end_cond = (input_tensor_ids == im_end_id) | (input_tensor_ids == slice_end_id)
            image_start_idx = torch.where(start_cond)[0] + 1
            image_end_idx = torch.where(end_cond)[0]
            valid_image_nums = min(len(image_start_idx), len(image_end_idx))
            image_bound = [
                torch.hstack([
                    image_start_idx[:valid_image_nums].unsqueeze(-1),
                    image_end_idx[:valid_image_nums].unsqueeze(-1),
                ])
            ]
        else:
            image_bound = [[]]

        # Step 5: Compute audio_bounds
        if inputs.audios:
            audio_start_idx = torch.where(input_tensor_ids == audio_start_id)[0]
            audio_end_idx = torch.where(input_tensor_ids == audio_end_id)[0]
            assert len(audio_start_idx) == len(audio_end_idx)
            audio_bounds = [torch.hstack([
                (audio_start_idx + 1).unsqueeze(-1),
                audio_end_idx.unsqueeze(-1),
            ])]
        else:
            audio_bounds = [[]]

        encoded = {
            'input_ids': input_ids,
            'labels': labels,
            'loss_scale': loss_scale,
            'image_bound': image_bound,
            'pixel_values': image_inputs['pixel_values'],
            'tgt_sizes': image_inputs['tgt_sizes'],
            'audio_features': audio_features,
            'audio_feature_lens': audio_feature_lens,
            'audio_bounds': audio_bounds,
        }
        return encoded

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        # Vision embeddings
        inputs_embeds, _ = model.get_vllm_embedding(inputs)
        # Audio embeddings — scatter audio features into the embedding
        inputs_embeds = model.get_omni_embedding(
            inputs,
            input_embeddings=inputs_embeds,
            chunk_length=getattr(self.config, 'audio_chunk_length', 1.0),
        )
        return {'inputs_embeds': inputs_embeds}

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = {}
        # Vision data
        for k in ['pixel_values', 'image_bound', 'tgt_sizes']:
            res[k] = self.gather_list(batch, k)

        # Audio data — collect from all samples
        all_audio_feats = []  # list of [N_i, 80, frames_i] tensors
        audio_feature_lens_list = []
        audio_bounds_list = []

        for b in batch:
            af = b.pop('audio_features', None)
            afl = b.pop('audio_feature_lens', None)
            ab = b.pop('audio_bounds', None)

            if af is not None and isinstance(af, torch.Tensor):
                all_audio_feats.append(af)
            if afl is not None:
                audio_feature_lens_list.extend(afl)
            if ab is not None:
                audio_bounds_list.extend(ab)

        # Re-pad audio features across the batch to the same max frame length
        if all_audio_feats:
            # Unpack per-sample tensors into individual segments, then re-pad
            segments = []
            for af in all_audio_feats:
                for i in range(af.shape[0]):
                    segments.append(af[i])  # [80, frames_i]
            res['audio_features'] = torch.nn.utils.rnn.pad_sequence(
                [s.transpose(0, 1) for s in segments],
                batch_first=True,
                padding_value=0.0,
            ).transpose(1, 2)  # [total_segments, 80, max_frames]
        else:
            res['audio_features'] = []

        res['audio_feature_lens'] = audio_feature_lens_list if audio_feature_lens_list else []
        res['audio_bounds'] = audio_bounds_list if audio_bounds_list else []

        res.update(Template._data_collator(self, batch, padding_to=padding_to))
        return res


register_template(
    Qwen3MixedTemplateMeta(
        MLLMTemplateType.minicpmo4_5,
        template_cls=MiniCPMO4_5Template,
        is_thinking=True,
    ))


class MiniCPMV4_6Template(Template):
    support_padding_free = True
    placeholder_tokens = ['<|image_pad|>', '<|video_pad|>']

    def init_env_args(self):
        super().init_env_args()
        self.downsample_mode = get_env_args('downsample_mode', str, '16x')
        self.max_slice_nums = get_env_args('max_slice_nums', int, 9)
        self.video_max_slice_nums = get_env_args('video_max_slice_nums', int, 1)
        self.max_num_frames = get_env_args('max_num_frames', int, 128)
        self.stack_frames = get_env_args('stack_frames', int, 1)
        self.transformers_version = version.parse(transformers.__version__)
        self.transformers_5_9 = self.transformers_version >= version.parse('5.9.0')

    def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None:
        super()._preprocess_inputs(inputs)
        # Inject downsample_mode into mm_processor_kwargs so that vLLM rollout
        # receives the correct mode via _encode_truncated -> _add_request.
        inputs.mm_processor_kwargs['downsample_mode'] = self.downsample_mode

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        if media_type == 'image':
            return ['<|image_pad|>\n']
        else:
            return ['<|video_pad|>\n']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        split_token = self._tokenize(self.tokenizer.eos_token)
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        for media_type in ['image', 'video']:
            mm_data = getattr(inputs, f'{media_type}s')
            media_token = f'<|{media_type}_pad|>'
            media_token_id = self._tokenize(media_token)[0]
            max_slice_nums = self.max_slice_nums if media_type == 'image' else self.video_max_slice_nums
            if mm_data:
                media_inputs = self.processor(
                    text=self.tokenizer.eos_token.join([media_token] * len(mm_data)),
                    images=inputs.images or None,
                    videos=inputs.videos or None,
                    return_tensors='pt',
                    add_special_tokens=False,
                    downsample_mode=self.downsample_mode,
                    stack_frames=self.stack_frames,
                    max_num_frames=self.max_num_frames,
                    max_slice_nums=max_slice_nums,
                )
                splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token)
                idx_list = findall(input_ids, media_token_id)

                def _get_new_tokens(i):
                    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 _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        res = {}
        pixel_values = [b['pixel_values'] for b in batch if b.get('pixel_values') is not None]
        if len(pixel_values) > 0:
            res['pixel_values'] = torch.concat(pixel_values, dim=-1)
        pixel_values_videos = [b['pixel_values_videos'] for b in batch if b.get('pixel_values_videos') is not None]
        if len(pixel_values_videos) > 0:
            res['pixel_values_videos'] = torch.concat(pixel_values_videos, dim=-1)

        for key in ['target_sizes', 'target_sizes_videos']:
            value = self.concat_tensor(batch, key, dim=0)
            if value is not None:
                res[key] = value

        # Inject downsample_mode so the model forward uses the same mode
        # as data preprocessing, keeping image token/feature counts aligned.
        res['downsample_mode'] = self.downsample_mode
        return res

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = super()._data_collator(batch, padding_to=padding_to)
        if self.padding_free:
            res.update(get_packed_seq_params(res['position_ids']))
        return res

    def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
        if self.padding_free and self.sequence_parallel_size <= 1 and not self.transformers_5_9:
            raise RuntimeError('MiniCPM-V 4.6 packing/padding_free with sequence_parallel_size=1 requires '
                               f'transformers>=5.9.0 (current: {self.transformers_version}). ')
        return super()._post_encode(model, inputs)


register_template(
    ChatmlTemplateMeta(
        MLLMTemplateType.minicpmv4_6,
        template_cls=MiniCPMV4_6Template,
        is_thinking=True,
        thinking_prefix='<think>\n',
        non_thinking_prefix='<think>\n\n</think>\n\n',
    ))

register_template(
    ChatmlTemplateMeta(
        LLMTemplateType.minicpm5,
        is_thinking=True,
        thinking_prefix='<think>\n',
        non_thinking_prefix='<think>\n\n</think>\n\n',
        agent_template='minicpm5',
    ))