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

from swift.utils import get_env_args, get_logger
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

logger = get_logger()


@dataclass
class MinimaxTemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=list)
    prompt: Prompt = field(default_factory=lambda: [
        '<beginning_of_sentence>user name=user\n{{QUERY}}<end_of_sentence>\n'
        '<beginning_of_sentence>ai name=assistant\n'
    ])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<end_of_sentence>\n'])
    suffix: Prompt = field(default_factory=lambda: ['<end_of_sentence>'])
    system_prefix: Optional[Prompt] = field(
        default_factory=lambda: ['<beginning_of_sentence>system ai_setting=assistant\n{{SYSTEM}}<end_of_sentence>\n'])


register_template(MinimaxTemplateMeta(LLMTemplateType.minimax))

register_template(
    MinimaxTemplateMeta(
        LLMTemplateType.minimax_m1,
        prefix=['<begin_of_document>'],
        system_prefix=[
            '<begin_of_document><beginning_of_sentence>system ai_setting=assistant\n{{SYSTEM}}<end_of_sentence>\n'
        ],
    ))


class MinimaxVLTemplate(Template):
    image_placeholder = ['<image>']
    skip_prompt = True

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type == 'image'
        return self.image_placeholder * inputs.all_image_tokens[index]

    def calc_num_image_tokens(self, image_inputs):
        from transformers.image_utils import get_image_size, to_numpy_array
        pixel_values = image_inputs['pixel_values']
        image_sizes = image_inputs['image_sizes']
        all_image_tokens = []
        if not image_inputs:
            return all_image_tokens

        if self.processor.process_image_mode == 'anyres':
            for pixel_value, image_size in zip(pixel_values, image_sizes):
                height, width = image_size
                num_image_tokens = self.processor.get_num_token(height, width, self.processor.grid_pinpoints,
                                                                self.processor.patch_size)
                all_image_tokens.append(num_image_tokens)
        elif self.processor.process_image_mode == 'resize':
            pixel_values = image_inputs['pixel_values']
            all_image_tokens = []
            for pixel_value in pixel_values:
                height, width = get_image_size(to_numpy_array(pixel_value))
                all_image_tokens.append(int(height * width / self.processor.patch_size**2))
        else:
            if self.processor.patch_size is not None:
                pixel_values = image_inputs['pixel_values']
                all_image_tokens = []
                for pixel_value in pixel_values:
                    height, width = get_image_size(to_numpy_array(pixel_value))
                    new_width, new_height = self.processor.get_hw_multiple_of(
                        (width, height), self.processor.patch_size, self.processor.max_size)
                    num_image_tokens = ((new_height // self.processor.patch_size) *
                                        (new_width // self.processor.patch_size))  # + 1
                    all_image_tokens.append(num_image_tokens)
            else:
                logger.warning_once(
                    'Expanding inputs for image tokens in MiniMaxVL01 should be done in processing. '
                    "Please add `patch_size` and `vision_feature_select_strategy` to the model's "
                    'processing config or set directly '
                    'with `processor.patch_size = {{patch_size}}` and processor.vision_feature_select_strategy = '
                    '{{vision_feature_select_strategy}}`. '
                    'Using processors without these attributes in the config is deprecated '
                    'and will throw an error in v4.47.')
                raise ValueError(
                    "You need to provide `patch_size` and `vision_feature_select_strategy` in the model's processing "
                    'config to expand inputs for image tokens.')
        return all_image_tokens

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        output_kwargs = self.processor._merge_kwargs(
            self.processor.MiniMaxVL01ProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
        )
        if inputs.images:
            image_inputs = self.processor.image_processor(
                inputs.images, **output_kwargs['images_kwargs'], return_tensors='pt')
            inputs.all_image_tokens = self.calc_num_image_tokens(image_inputs)
        else:
            image_inputs = {}
        encoded = super()._encode(inputs)
        for key in image_inputs:
            encoded[key] = image_inputs[key]
        return encoded

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


register_template(MinimaxTemplateMeta(LLMTemplateType.minimax_vl, template_cls=MinimaxVLTemplate))


@dataclass
class MinimaxM2TemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: [']~!b[]~b]system\n{{SYSTEM}}[e~[\n'])
    prompt: Prompt = field(default_factory=lambda: [']~b]user\n{{QUERY}}[e~[\n]~b]ai\n'])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['[e~[\n'])
    suffix: Prompt = field(default_factory=lambda: ['[e~[\n'])
    agent_template: Optional[str] = 'minimax_m2'
    is_thinking: bool = True
    thinking_prefix: str = '<think>\n'


register_template(
    MinimaxM2TemplateMeta(
        LLMTemplateType.minimax_m2,
        default_system='You are MiniMax-M2, a helpful AI assistant built by MiniMax. Knowledge cutoff: 2025-06.',
    ))

register_template(
    MinimaxM2TemplateMeta(
        LLMTemplateType.minimax_m2_1,
        default_system='You are a helpful assistant. Your name is MiniMax-M2.1 and is built by MiniMax.',
    ))

register_template(
    MinimaxM2TemplateMeta(
        LLMTemplateType.minimax_m2_5,
        default_system='You are a helpful assistant. Your name is MiniMax-M2.5 and is built by MiniMax.',
    ))

register_template(
    MinimaxM2TemplateMeta(
        LLMTemplateType.minimax_m2_7,
        default_system='You are a helpful assistant. Your name is MiniMax-M2.7 and is built by MiniMax.',
    ))

_MINIMAX_M3_IDENTITY = ('Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. '
                        'Founded in early 2022, MiniMax is a global AI foundation model company committed to '
                        'advancing the frontiers of AI towards AGI.')

_MINIMAX_M3_THINKING_BASE = (
    'You have a thinking capability that allows you to reason step by step before responding. '
    'When thinking is enabled, wrap your reasoning in <mm:think></mm:think> tags before your '
    'response. When thinking is disabled, begin your response directly after the </mm:think> '
    'prefix. When thinking is adaptive, decide on your own whether to think for the current turn.')

_MINIMAX_M3_THINKING_MODE_TEXT = {
    'enabled': ('Current thinking mode: enabled. You MUST think step by step before every response, '
                'including after receiving function/tool results.'),
    'disabled':
    'Current thinking mode: disabled. Do not output any thinking process.',
    'adaptive': ('Current thinking mode: adaptive. You are encouraged to think for complex '
                 'decision-making, multi-step reasoning, or when analyzing function/tool results.'),
}

_MINIMAX_M3_DEFAULT_DEVELOPER = 'You are a helpful assistant.'


def _build_m3_system_block(thinking_mode: str = 'adaptive') -> str:
    mode_text = _MINIMAX_M3_THINKING_MODE_TEXT.get(thinking_mode, _MINIMAX_M3_THINKING_MODE_TEXT['adaptive'])
    return (f'{_MINIMAX_M3_IDENTITY}'
            f'\n\n<thinking_instructions>\n{_MINIMAX_M3_THINKING_BASE}\n{mode_text}\n</thinking_instructions>')


class MinimaxM3VLTemplate(Template):
    image_token = ']<]image[>['
    video_token = ']<]video[>['
    placeholder_tokens = [']<]image[>[', ']<]video[>[']

    def init_env_args(self):
        super().init_env_args()
        # thinking_mode: "enabled" / "disabled" / "adaptive"
        self.thinking_mode = get_env_args('thinking_mode', str, 'disabled')
        self.chat_template_kwargs['thinking_mode'] = self.thinking_mode
        # Map thinking_mode to enable_thinking for the broader framework
        if self.thinking_mode == 'disabled':
            self.enable_thinking = False
        else:
            self.enable_thinking = True

    def _get_thinking_mode(self, inputs=None) -> str:
        thinking_mode = None if inputs is None else inputs.chat_template_kwargs.get('thinking_mode')
        if thinking_mode is None:
            thinking_mode = self.chat_template_kwargs.get('thinking_mode', 'adaptive')
        return thinking_mode

    def _get_enable_thinking(self, inputs=None):
        thinking_mode = self._get_thinking_mode(inputs)
        return thinking_mode != 'disabled'

    def _get_response_prefix(self, inputs=None):
        # Check explicit override first
        response_prefix = None if inputs is None else inputs.chat_template_kwargs.get('response_prefix')
        if response_prefix is not None:
            return response_prefix
        if self.response_prefix is not None:
            return self.response_prefix
        thinking_mode = self._get_thinking_mode(inputs)
        if thinking_mode == 'enabled':
            return self.template_meta.thinking_prefix  # '<mm:think>'
        elif thinking_mode == 'disabled':
            return self.template_meta.non_thinking_prefix  # '</mm:think>'
        else:  # adaptive
            return ''  # No prefix, let model decide

    def _get_system(self, inputs: StdTemplateInputs) -> str:
        system = super()._get_system(inputs)
        thinking_mode = self._get_thinking_mode(inputs)
        system_block = _build_m3_system_block(thinking_mode)
        return f'{system_block}[e~[\n]~b]developer\n{system or ""}'

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        if media_type == 'image':
            return [self.image_token]
        elif media_type == 'video':
            return [self.video_token]
        else:
            raise ValueError(f'Unsupported media type for MiniMax-M3 VL: {media_type}')

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        if not inputs.images and not inputs.videos:
            return encoded
        media_text_parts = ([self.image_token] * len(inputs.images) + [self.video_token] * len(inputs.videos))
        media_inputs = self.processor(
            text=self.tokenizer.eos_token.join(media_text_parts),
            images=inputs.images or None,
            videos=inputs.videos or None,
            return_tensors='pt',
        )
        split_token = self._tokenize(self.tokenizer.eos_token)
        splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token)
        media_inputs.pop('input_ids', None)
        media_inputs.pop('attention_mask', None)

        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)

        idx_list = []
        for key in ['image', 'video']:
            token_id = getattr(self.config, f'{key}_token_id', None)
            if token_id is None:
                continue
            idx_list += findall(input_ids, token_id)
        sorted_order = sorted(range(len(idx_list)), key=lambda i: idx_list[i])
        idx_list = [idx_list[i] for i in sorted_order]
        splited_tokens = [splited_tokens[i] for i in sorted_order]

        def _get_new_tokens(i):
            return splited_tokens[i]

        if idx_list:
            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


@dataclass
class MinimaxM3VLTemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: [']~!b[]~b]system\n{{SYSTEM}}[e~[\n'])
    prompt: Prompt = field(default_factory=lambda: [']~b]user\n{{QUERY}}[e~[\n]~b]ai\n'])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['[e~[\n'])
    suffix: Prompt = field(default_factory=lambda: ['[e~[\n'])
    default_system: Optional[str] = _MINIMAX_M3_DEFAULT_DEVELOPER
    agent_template: Optional[str] = 'minimax_m3'
    is_thinking: bool = True
    thinking_prefix: str = '<mm:think>'
    non_thinking_prefix: str = '</mm:think>'
    history_thinking_prefix: str = '</mm:think>'


register_template(MinimaxM3VLTemplateMeta(MLLMTemplateType.minimax_m3_vl, template_cls=MinimaxM3VLTemplate))