"""
MLX-based KimiVL Processor.

This module provides an MLX-native processor for KimiVL models that:
1. Uses a pre-converted fast tokenizer (no tiktoken dependency)
2. Provides an MLX-based image processor (no torch/torchvision dependency)
3. Patches missing functions for transformers 5.0 compatibility
"""

import json
import math
import warnings
from pathlib import Path
from typing import List, Optional, Tuple, Union

import mlx.core as mx
import transformers.processing_utils as processing_utils
from PIL import Image
from transformers import AutoTokenizer
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor
from transformers.image_utils import ImageInput, make_list_of_images, valid_images
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import TensorType

from .config import ModelConfig


def _validate_images_text_input_order(images, text):
    """
    Validate and potentially swap the order of images and text arguments.

    This function checks if the arguments are in the correct order (images first, text second)
    for backward compatibility. If text is passed as the first argument and images as the second,
    it swaps them and issues a deprecation warning.

    Args:
        images: The images argument (should be image-like objects or None)
        text: The text argument (should be strings or None)

    Returns:
        Tuple of (images, text) in the correct order
    """
    # Check if arguments are swapped (text passed as images, images passed as text)
    if images is not None and text is not None:
        # If 'images' looks like text and 'text' looks like images, swap them
        images_is_text = isinstance(images, str) or (
            isinstance(images, (list, tuple))
            and len(images) > 0
            and isinstance(images[0], str)
        )
        text_is_image = not isinstance(text, str) and not (
            isinstance(text, (list, tuple))
            and len(text) > 0
            and isinstance(text[0], str)
        )

        if images_is_text and text_is_image:
            warnings.warn(
                "You passed text as the first argument and images as the second. "
                "This is deprecated and will be removed in a future version. "
                "Please pass images first and text second.",
                FutureWarning,
            )
            return text, images

    return images, text


# Add the function to transformers.processing_utils if it doesn't exist
if not hasattr(processing_utils, "_validate_images_text_input_order"):
    processing_utils._validate_images_text_input_order = (
        _validate_images_text_input_order
    )

# Also add Unpack if it doesn't exist (for older Python versions)
if not hasattr(processing_utils, "Unpack"):
    try:
        from typing import Unpack

        processing_utils.Unpack = Unpack
    except ImportError:
        from typing_extensions import Unpack

        processing_utils.Unpack = Unpack


def _ensure_gpt2_bytes_to_unicode():
    """
    Compatibility shim for tokenizer remote code expecting legacy GPT-2 helper.

    Newer Transformers versions removed bytes_to_unicode from
    transformers.models.gpt2.tokenization_gpt2, but remote tokenizers may still
    import it from there.
    """
    try:
        import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
    except Exception:
        return

    if hasattr(gpt2_tokenization, "bytes_to_unicode"):
        return

    def bytes_to_unicode():
        bs = (
            list(range(ord("!"), ord("~") + 1))
            + list(range(ord("¡"), ord("¬") + 1))
            + list(range(ord("®"), ord("ÿ") + 1))
        )
        cs = bs[:]
        n = 0
        for b in range(2**8):
            if b not in bs:
                bs.append(b)
                cs.append(2**8 + n)
                n += 1
        cs = [chr(n) for n in cs]
        return dict(zip(bs, cs))

    gpt2_tokenization.bytes_to_unicode = bytes_to_unicode


# CLIP-style normalization constants
OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)


class KimiVLImageProcessor(BaseImageProcessor):

    model_input_names = ["pixel_values", "image_grid_hws"]

    def __init__(
        self,
        patch_size: int = 14,
        pad_input: bool = False,
        image_mean: Tuple[float, float, float] = OPENAI_DATASET_MEAN,
        image_std: Tuple[float, float, float] = OPENAI_DATASET_STD,
        in_token_limit: int = 4096,
        merge_kernel_size: List[int] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.in_token_limit = in_token_limit
        self.patch_size = patch_size
        self.pad_input = pad_input
        self.image_mean = image_mean
        self.image_std = image_std
        self.merge_kernel_size = (
            merge_kernel_size if merge_kernel_size is not None else [2, 2]
        )

    def rescale(
        self, image: Image.Image, merge_kernel_size: List[int] = None
    ) -> Image.Image:
        """Rescale image to fit within token limits and pad/crop to patch boundaries."""
        if merge_kernel_size is None:
            merge_kernel_size = self.merge_kernel_size

        w, h = image.size
        patch_size = self.patch_size

        # Rescale if exceeds token limit
        if (w // patch_size) * (h // patch_size) > self.in_token_limit:
            scale = math.sqrt(
                self.in_token_limit / ((w // patch_size) * (h // patch_size))
            )
            new_w, new_h = int(w * scale), int(h * scale)
            image = image.resize((new_w, new_h), Image.Resampling.BICUBIC)

        if self.pad_input:
            new_w, new_h = image.size
            pad_size_h = merge_kernel_size[0] * patch_size
            pad_size_w = merge_kernel_size[1] * patch_size

            pad_h = (pad_size_h - new_h % pad_size_h) % pad_size_h
            pad_w = (pad_size_w - new_w % pad_size_w) % pad_size_w

            if pad_h > 0 or pad_w > 0:
                # Pad image (bottom and right padding)
                new_image = Image.new(
                    image.mode, (new_w + pad_w, new_h + pad_h), (0, 0, 0)
                )
                new_image.paste(image, (0, 0))
                image = new_image
        else:
            new_w, new_h = image.size
            # Ensure dimensions are divisible by merge_kernel_size * patch_size
            # so that the grid dimensions are divisible by merge_kernel_size
            crop_size_w = merge_kernel_size[1] * patch_size
            crop_size_h = merge_kernel_size[0] * patch_size
            new_w = new_w - new_w % crop_size_w
            new_h = new_h - new_h % crop_size_h
            # Center crop
            left = (image.size[0] - new_w) // 2
            top = (image.size[1] - new_h) // 2
            image = image.crop((left, top, left + new_w, top + new_h))

        w, h = image.size
        if w // patch_size >= 512 or h // patch_size >= 512:
            raise ValueError("Exceed pos emb")

        return image

    def to_mlx(self, image: Image.Image) -> mx.array:
        """Convert PIL image to MLX array in CHW format, normalized to [0, 1]."""
        image = image.convert("RGB")
        w, h = image.size
        # Convert PIL image to MLX array directly via bytes
        arr = mx.array(list(image.getdata()), dtype=mx.float32).reshape(h, w, 3) / 255.0
        # Convert from HWC to CHW format
        arr = arr.transpose(2, 0, 1)
        return arr

    def normalize(self, image: mx.array) -> mx.array:
        """Normalize image with CLIP-style mean and std."""
        mean = mx.array(self.image_mean, dtype=mx.float32).reshape(3, 1, 1)
        std = mx.array(self.image_std, dtype=mx.float32).reshape(3, 1, 1)
        return (image - mean) / std

    def patchify(self, image: mx.array) -> Tuple[mx.array, Tuple[int, int]]:
        """Convert image to patches."""
        patch_size = self.patch_size
        C, H, W = image.shape

        # Reshape to (C, H//p, p, W//p, p) then to (num_patches, C, p, p)
        patches = image.reshape(
            C, H // patch_size, patch_size, W // patch_size, patch_size
        )
        # Permute to (H//p, W//p, C, p, p)
        patches = patches.transpose(1, 3, 0, 2, 4)
        # Flatten to (num_patches, C, p, p)
        patches = patches.reshape(-1, C, patch_size, patch_size)

        grid_hw = (H // patch_size, W // patch_size)
        return patches, grid_hw

    def _preprocess(self, image: ImageInput) -> Tuple[mx.array, Tuple[int, int]]:
        """
        Preprocess image and patchify it.

        Args:
            image: Image to preprocess.

        Returns:
            patches: mx.array
            grid_hw: Tuple[int, int]
        """
        image = self.rescale(image, self.merge_kernel_size)
        image = self.to_mlx(image)
        image = self.normalize(image)
        patches, grid_hw = self.patchify(image)
        return patches, grid_hw

    def preprocess(
        self,
        images: ImageInput,
        return_tensors: Optional[Union[str, TensorType]] = None,
        **kwargs,
    ) -> BatchFeature:
        """Process images and return BatchFeature."""
        images = make_list_of_images(images)

        if not valid_images(images):
            raise ValueError(
                "Invalid image type. Must be of type PIL.Image.Image or mx.array."
            )

        pixel_values_list = []
        image_grid_hws = []

        for image in images:
            # Convert MLX arrays to PIL Images if needed
            if isinstance(image, mx.array):
                # Ensure we're working with the array values
                arr = image
                if arr.ndim == 3 and arr.shape[0] in [1, 3, 4]:
                    # CHW format, convert to HWC
                    arr = arr.transpose(1, 2, 0)
                # Convert to uint8 for PIL
                if arr.dtype in [mx.float32, mx.float16, mx.bfloat16]:
                    arr = (arr * 255).astype(mx.uint8)
                # Convert to PIL via list (MLX -> list -> PIL)
                h, w, _ = arr.shape
                flat_data = arr.reshape(-1).tolist()
                image = Image.frombytes("RGB", (w, h), bytes(flat_data))

            patches, image_grid_hw = self._preprocess(image)
            pixel_values_list.append(patches)
            image_grid_hws.append(image_grid_hw)

        pixel_values = mx.concatenate(pixel_values_list, axis=0)
        image_grid_hws = mx.array(image_grid_hws)

        # Return MLX arrays directly
        data = {
            "pixel_values": pixel_values,
            "image_grid_hws": image_grid_hws,
        }

        return BatchFeature(data=data, tensor_type=return_tensors)

    def __call__(
        self,
        images: ImageInput,
        return_tensors: Optional[Union[str, TensorType]] = None,
        **kwargs,
    ) -> BatchFeature:
        """Make the image processor callable."""
        return self.preprocess(images, return_tensors=return_tensors, **kwargs)


class KimiVLProcessor(ProcessorMixin):
    """
    MLX-based processor for KimiVL that doesn't require torch/torchvision.

    Constructs a KimiVL processor which wraps a KimiVL image processor and a tokenizer
    into a single processor.
    """

    attributes = ["image_processor", "tokenizer"]
    valid_kwargs = ["chat_template"]
    image_processor_class = "KimiVLImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        chat_template=None,
        **kwargs,
    ):
        self.image_token = "<|media_pad|>"
        if image_processor is None:
            image_processor = KimiVLImageProcessor()
        super().__init__(image_processor, tokenizer, chat_template=chat_template)

    def __call__(
        self,
        images: ImageInput = None,
        text: Union[
            TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]
        ] = None,
        **kwargs,
    ) -> BatchFeature:
        """
        Main method to prepare for the model one or several sequences(s) and image(s).

        Args:
            images: The image or batch of images to be prepared.
            text: The sequence or batch of sequences to be encoded.
            return_tensors: If set, will return tensors of a particular framework.

        Returns:
            BatchFeature with input_ids, attention_mask, and pixel_values.
        """
        if images is None and text is None:
            raise ValueError("You have to specify at least one of `images` or `text`.")

        # Check if images and text inputs are reversed for BC
        images, text = _validate_images_text_input_order(images, text)

        # Extract return_tensors from kwargs (unused, we always return MLX arrays)
        kwargs.pop("return_tensors", None)

        # Process images
        if images is not None:
            image_inputs = self.image_processor(images)
            image_grid_hws = image_inputs["image_grid_hws"]
        else:
            image_inputs = {}
            image_grid_hws = None

        # Process text
        if isinstance(text, str):
            text = [text]
        elif text is not None and not isinstance(text, list):
            raise ValueError(
                "Invalid input text. Please provide a string, or a list of strings"
            )

        # Replace image tokens with the correct number of placeholder tokens
        if image_grid_hws is not None and text is not None:
            merge_length = (
                self.image_processor.merge_kernel_size[0]
                * self.image_processor.merge_kernel_size[1]
            )
            index = 0
            for i in range(len(text)):
                while self.image_token in text[i]:
                    # Use mx.prod for MLX arrays
                    grid_hw = image_grid_hws[index]
                    num_placeholders = int(mx.prod(grid_hw).item()) // merge_length
                    text[i] = text[i].replace(
                        self.image_token,
                        "<|placeholder|>" * num_placeholders,
                        1,
                    )
                    index += 1
                text[i] = text[i].replace("<|placeholder|>", self.image_token)

        # Tokenize text
        # Note: The TikToken tokenizer doesn't work properly with transformers' standard
        # __call__ method due to issues with the pad function. We use encode() directly.
        if text is not None:
            # Encode each text and build the result manually
            all_input_ids = []
            for t in text:
                ids = self.tokenizer.encode(t)
                all_input_ids.append(ids)

            # Pad sequences to the same length if needed
            max_len = max(len(ids) for ids in all_input_ids)
            pad_token_id = self.tokenizer.pad_token_id or 0

            padded_input_ids = []
            attention_masks = []
            for ids in all_input_ids:
                padding_length = max_len - len(ids)
                padded_ids = ids + [pad_token_id] * padding_length
                mask = [1] * len(ids) + [0] * padding_length
                padded_input_ids.append(padded_ids)
                attention_masks.append(mask)

            # Convert to MLX arrays
            text_inputs = {
                "input_ids": mx.array(padded_input_ids),
                "attention_mask": mx.array(attention_masks),
            }
        else:
            text_inputs = {}

        data = {**text_inputs, **image_inputs}
        image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
        if image_token_id is not None:
            data["image_token_id"] = int(image_token_id)
        return BatchFeature(data=data)

    def batch_decode(self, *args, **kwargs):
        """Forward to tokenizer's batch_decode."""
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """Forward to tokenizer's decode."""
        return self.tokenizer.decode(*args, **kwargs)

    def apply_chat_template(
        self,
        conversation,
        chat_template=None,
        add_generation_prompt=False,
        tokenize=False,
        **kwargs,
    ):
        """Apply chat template to the conversation."""
        # Use provided template, processor's template, or tokenizer's template
        if chat_template is None:
            chat_template = self.chat_template
        if chat_template is None:
            chat_template = getattr(self.tokenizer, "chat_template", None)
        if chat_template is None:
            raise ValueError(
                "No chat template found. Please provide a chat_template argument "
                "or ensure the tokenizer has a chat_template attribute."
            )

        # Use jinja2 to render the template
        try:
            from jinja2 import Template
        except ImportError:
            raise ImportError("jinja2 is required for apply_chat_template")

        template = Template(chat_template)
        rendered = template.render(
            messages=conversation,
            add_generation_prompt=add_generation_prompt,
            **kwargs,
        )

        if tokenize:
            return self.tokenizer.encode(rendered)
        return rendered

    @property
    def model_input_names(self):
        """Get the model input names from tokenizer and image processor."""
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        """Load the processor from a pretrained model path."""
        from huggingface_hub import hf_hub_download

        kwargs.pop("trust_remote_code", None)
        _ensure_gpt2_bytes_to_unicode()

        model_path = Path(pretrained_model_name_or_path)
        is_local = model_path.exists() and model_path.is_dir()
        tokenizer = AutoTokenizer.from_pretrained(
            str(model_path) if is_local else pretrained_model_name_or_path,
            trust_remote_code=True,
            local_files_only=is_local,
        )

        # Load image processor config and create our processor
        image_processor_config = {}
        try:
            if is_local:
                config_path = model_path / "config.json"
            else:
                config_path = Path(
                    hf_hub_download(pretrained_model_name_or_path, "config.json")
                )
            with open(config_path, "r", encoding="utf-8") as f:
                config_dict = json.load(f)
            config = ModelConfig.from_dict(config_dict)
            if hasattr(config, "vision_config"):
                vision_config = config.vision_config
                if hasattr(vision_config, "patch_size"):
                    image_processor_config["patch_size"] = vision_config.patch_size
                if hasattr(vision_config, "in_token_limit"):
                    image_processor_config["in_token_limit"] = (
                        vision_config.in_token_limit
                    )
                if hasattr(vision_config, "merge_kernel_size"):
                    image_processor_config["merge_kernel_size"] = (
                        vision_config.merge_kernel_size
                    )
        except Exception:
            pass

        image_processor = KimiVLImageProcessor(**image_processor_config)

        # Load chat template from jinja file if not already set on tokenizer
        chat_template = getattr(tokenizer, "chat_template", None)
        if chat_template is None:
            try:
                if is_local:
                    jinja_path = model_path / "chat_template.jinja"
                else:
                    jinja_path = Path(
                        hf_hub_download(
                            pretrained_model_name_or_path, "chat_template.jinja"
                        )
                    )
                if jinja_path.exists():
                    chat_template = jinja_path.read_text(encoding="utf-8")
                    # Set chat_template on tokenizer so apply_chat_template works
                    tokenizer.chat_template = chat_template
            except Exception:
                pass

        return cls(
            image_processor=image_processor,
            tokenizer=tokenizer,
            chat_template=chat_template,
        )


from transformers import AutoProcessor

_original_auto_processor_from_pretrained = AutoProcessor.from_pretrained


@classmethod
def _patched_auto_processor_from_pretrained(
    cls, pretrained_model_name_or_path, **kwargs
):
    """Patched from_pretrained that returns KimiVLProcessor for kimi_vl models."""
    from huggingface_hub import hf_hub_download

    model_path = Path(pretrained_model_name_or_path)
    is_local = model_path.exists() and model_path.is_dir()

    # Check if this is a kimi_vl model
    is_kimi_vl = False
    try:
        if is_local:
            config_path = model_path / "config.json"
        else:
            config_path = Path(
                hf_hub_download(pretrained_model_name_or_path, "config.json")
            )
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        is_kimi_vl = config.get("model_type", "").lower() == "kimi_vl"
    except Exception:
        pass

    if is_kimi_vl:
        return KimiVLProcessor.from_pretrained(pretrained_model_name_or_path, **kwargs)

    return _original_auto_processor_from_pretrained.__func__(
        cls, pretrained_model_name_or_path, **kwargs
    )


AutoProcessor.from_pretrained = _patched_auto_processor_from_pretrained
