import glob
import importlib
import inspect
import json
import logging
from io import BytesIO
from pathlib import Path
from textwrap import dedent
from typing import Any, Dict, List, Optional, Tuple, Union

import mlx.core as mx
import mlx.nn as nn
import numpy as np
import requests
from huggingface_hub import snapshot_download
from mlx.utils import tree_flatten, tree_map
from PIL import Image, ImageOps
from transformers import AutoProcessor
from transformers.processing_utils import ProcessorMixin

from .models.base import BaseImageProcessor
from .tokenizer_utils import load_tokenizer
from .trainer.utils import apply_lora_layers

# Modes that support activation quantization
ACTIVATION_QUANTIZATION_MODES = {"nvfp4", "mxfp8"}

# Constants
MODEL_REMAPPING = {
    "llava_qwen2": "fastvlm",  # Apple's FastVLM, note it's different to the one below
    "llava-qwen2": "llava_bunny",
    "bunny-llama": "llava_bunny",
    "lfm2-vl": "lfm2_vl",
    "cohere2_vision": "aya_vision",
    "jvlm": "jina_vlm",
    "phi4-siglip": "phi4_siglip",
    "sam3_video": "sam3",
}

MAX_FILE_SIZE_GB = 5

MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]


def quantize_activations(model: nn.Module) -> nn.Module:

    def _maybe_qq(m: nn.Module) -> nn.Module:
        """Convert a QuantizedLinear layer to QQLinear if compatible."""
        if isinstance(m, nn.QuantizedLinear):
            if m.mode not in ACTIVATION_QUANTIZATION_MODES:
                raise ValueError(
                    f"Mode ({m.mode}) does not support activation quantization. "
                    f"Supported modes: {', '.join(ACTIVATION_QUANTIZATION_MODES)}"
                )
            if m.get("bias", False):
                raise ValueError(
                    "Linear layer with bias does not support activation quantization"
                )
            # Get dimensions from quantized weight
            out_dims, in_dims = m.weight.shape
            in_dims *= 32 // m.bits
            qq = nn.QQLinear(in_dims, out_dims, m.group_size, m.bits, m.mode)
            qq.quantize()
            return qq
        return m

    leaves = tree_map(_maybe_qq, model.leaf_modules(), is_leaf=nn.Module.is_module)
    model.update_modules(leaves)

    return model


def skip_multimodal_module(path: str) -> bool:
    """
    Check if a multimodal module (vision/audio) should skip quantization.

    Args:
        path: The module path to check

    Returns:
        bool: True if the module is multimodal and should skip quantization, False otherwise
    """
    return (
        "vision_model" in path
        or "vision_tower" in path
        or "vl_connector" in path
        or "sam_model" in path
        or "audio_model" in path
        or "audio_tower" in path
        or "code_predictor" in path
    )


def get_model_and_args(config: dict):
    """
    Retrieve the model object based on the configuration.

    Args:
        config (dict): The model configuration.

    Returns:
        A tuple containing the Model class and the ModelArgs class.
    """
    model_type = config["model_type"].lower()

    model_type = MODEL_REMAPPING.get(model_type, model_type)

    try:
        arch = importlib.import_module(f"mlx_vlm.models.{model_type}")
    except ImportError as e:
        msg = f"Model type {model_type} not supported. Error: {e}"
        logging.error(msg)
        raise ValueError(msg)

    return arch, model_type


def get_model_path(
    path_or_hf_repo: str, revision: Optional[str] = None, force_download: bool = False
) -> Path:
    """
    Ensures the model is available locally. If the path does not exist locally,
    it is downloaded from the Hugging Face Hub.

    Args:
        path_or_hf_repo (str): The local path or Hugging Face repository ID of the model.
        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.

    Returns:
        Path: The path to the model.
    """
    model_path = Path(path_or_hf_repo)
    if not model_path.exists():
        model_path = Path(
            snapshot_download(
                repo_id=path_or_hf_repo,
                revision=revision,
                allow_patterns=[
                    "*.json",
                    "*.safetensors",
                    "*.py",
                    "*.model",
                    "*.tiktoken",
                    "*.txt",
                    "*.jinja",
                ],
                force_download=force_download,
            )
        )
    return model_path


def load_model(model_path: Path, lazy: bool = False, **kwargs) -> nn.Module:
    """
    Load and initialize the model from a given path.

    Args:
        model_path (Path): The path to load the model from.
        lazy (bool): If False eval the model parameters to make sure they are
            loaded in memory before returning, otherwise they will be loaded
            when needed. Default: ``False``
        revision (str, optional): A revision id which can be a branch name,
            a tag, or a commit hash. Default: ``None``.
        quantize_activations (bool, optional): If True, convert QuantizedLinear layers
            to QQLinear layers for activation quantization. Only supported for models
            quantized with 'nvfp4' or 'mxfp8' modes. Default: ``False``.
        quantize_activations (bool, optional): If True, convert QuantizedLinear layers
            to QQLinear layers for activation quantization. Only supported for models
            quantized with 'nvfp4' or 'mxfp8' modes. Default: ``False``.

    Returns:
        nn.Module: The loaded and initialized model.

    Raises:
        FileNotFoundError: If the weight files (.safetensors) are not found.
        ValueError: If the model class or args class are not found or cannot be instantiated.
    """
    config = load_config(model_path, **kwargs)

    # Find all .safetensors files in the model_path, excluding consolidated model weights
    weight_files = [
        wf
        for wf in glob.glob(str(model_path / "*.safetensors"))
        if not wf.endswith("consolidated.safetensors")
    ]

    if not weight_files:
        logging.error(f"No safetensors found in {model_path}")
        message = f"""
No safetensors found in {model_path}
Create safetensors using the following code:
```
from transformers import AutoModelForCausalLM, AutoProcessor

model_id= "<huggingface_model_id>"
model = AutoModelForCausalLM.from_pretrained(model_id)
processor = AutoProcessor.from_pretrained(model_id)

model.save_pretrained("<local_dir>")
processor.save_pretrained("<local_dir>")
```
Then use the <local_dir> as the --hf-path in the convert script.
```
python -m mlx_vlm.convert --hf-path <local_dir> --mlx-path <mlx_dir>
```
        """
        raise FileNotFoundError(message)

    weights = {}
    for wf in weight_files:
        weights.update(mx.load(wf))

    import safetensors

    with safetensors.safe_open(weight_files[0], framework="np") as f:
        is_mlx_format = f.metadata() and f.metadata().get("format") == "mlx"

    model_class, _ = get_model_and_args(config=config)

    # Initialize text and vision configs if not present
    config.setdefault("text_config", config.pop("llm_config", {}))
    config.setdefault("vision_config", {})
    config.setdefault("audio_config", {})

    # Initialize model config and update it with module configs
    model_config = model_class.ModelConfig.from_dict(config)
    modules = ["text", "vision", "perceiver", "projector", "audio"]
    model_config = update_module_configs(model_config, model_class, config, modules)

    model = model_class.Model(model_config)

    if not is_mlx_format:
        # Sanitize weights
        weights = sanitize_weights(model, weights)

        if hasattr(model, "thinker") and hasattr(model.thinker, "sanitize"):
            weights = sanitize_weights(model.thinker, weights)
            weights = sanitize_weights(model.thinker.vision_tower, weights)
            weights = sanitize_weights(model.thinker.audio_tower, weights)
            weights = sanitize_weights(model.thinker.language_model, weights)
            weights = sanitize_weights(model.code2wav, weights)
            weights = sanitize_weights(model.talker, weights)
        else:
            weights = sanitize_weights(
                model_class.VisionModel, weights, model_config.vision_config
            )
            weights = sanitize_weights(
                model_class.LanguageModel, weights, model_config.text_config
            )
            if hasattr(model_class, "AudioModel"):
                weights = sanitize_weights(
                    model_class.AudioModel, weights, model_config.audio_config
                )

    if not "quantization" in config:
        quantization_config = config.get("quantization_config", None)
        if quantization_config is None:
            text_config = config.get("text_config", {})
            quantization_config = text_config.get("quantization_config", None)
            if quantization_config is not None:
                config["quantization_config"] = quantization_config

        if quantization_config is not None:
            quant_method = quantization_config.get("quant_method")
            quantization = None
            if quant_method == "compressed-tensors":
                quantization = {"group_size": 32, "bits": 4, "mode": "affine"}
            elif quant_method == "mxfp4":
                quantization = {"group_size": 32, "bits": 4, "mode": "mxfp4"}
            elif quant_method in ("awq", "gptq", "bitnet"):
                logging.warning(
                    "Quantization method %s is not supported in mlx_vlm.load_model()",
                    quant_method,
                )

            if quantization is not None:
                config["quantization"] = quantization
                config["quantization_config"] = quantization

    if (quantization := config.get("quantization", None)) is not None:
        # Handle legacy models which may or may not have vision quantized
        # TODO: Re-upload the models with the new quantization config and remove this
        skip_vision = config.get("vision_config", {}).get("skip_vision", False)

        def get_class_predicate(p, m):
            # Always skip vision and audio models
            if skip_multimodal_module(p) and skip_vision:
                return False
            # Handle custom per layer quantizations
            if p in config["quantization"]:
                return config["quantization"][p]
            if not hasattr(m, "to_quantized"):
                return False
            # Skip layers not divisible by 64
            if hasattr(m, "weight") and m.weight.size % 64 != 0:
                return False
            # Handle legacy models which may not have everything quantized
            return f"{p}.scales" in weights

        nn.quantize(
            model,
            group_size=quantization["group_size"],
            bits=quantization["bits"],
            mode=quantization.get("mode", "affine"),
            class_predicate=get_class_predicate,
        )

    if kwargs.get("quantize_activations", False):
        if quantization is None:
            raise ValueError(
                "Activation quantization requires the model to be quantized first. "
                "Please use a quantized model with mode 'nvfp4' or 'mxfp8'."
            )
        model = quantize_activations(model)

    model.load_weights(list(weights.items()))

    if not lazy:
        mx.eval(model.parameters())

    model.eval()
    return model


def sanitize_weights(model_obj, weights, config=None):
    """Helper function to sanitize weights if the model has a sanitize method"""
    if hasattr(model_obj, "sanitize"):
        if config is not None:
            model_obj = model_obj(config)
        weights = model_obj.sanitize(weights)
    return weights


def update_module_configs(model_config, model_class, config, modules):
    """Updates configuration for model modules like text and vision modules.

    Args:
        model_config: The model configuration object that will be updated
        model_class: The model class containing component config classes
        config: Dictionary containing configuration parameters
        modules: List of module names to update configs for (e.g. ["text", "vision"])

    Returns:
        The updated model_config object
    """
    for config_name in modules:
        config_attr = f"{config_name}_config"
        if hasattr(model_config, config_attr):
            config_class = getattr(model_class, f"{config_name.title()}Config")
            setattr(
                model_config, config_attr, config_class.from_dict(config[config_attr])
            )
    return model_config


def load(
    path_or_hf_repo: str,
    adapter_path: Optional[str] = None,
    lazy: bool = False,
    revision: Optional[str] = None,
    **kwargs,
) -> Tuple[nn.Module, ProcessorMixin]:
    """
    Load the model and tokenizer from a given path or a huggingface repository.

    Args:
        path_or_hf_repo (Path): The path or the huggingface repository to load the model from.
        tokenizer_config (dict, optional): Configuration parameters specifically for the tokenizer.
            Defaults to an empty dictionary.
        adapter_path (str, optional): Path to the LoRA adapters. If provided, applies LoRA layers
            to the model. Default: ``None``.
        lazy (bool): If False eval the model parameters to make sure they are
            loaded in memory before returning, otherwise they will be loaded
            when needed. Default: ``False``
        revision (str, optional): A revision id which can be a branch name,
            a tag, or a commit hash. Default: ``None``.
        quantize_activations (bool, optional): If True, convert QuantizedLinear layers
            to QQLinear layers for activation quantization. Only supported for models
            quantized with 'nvfp4' or 'mxfp8' modes. Default: ``False``.

        quantize_activations (bool, optional): If True, convert QuantizedLinear layers
            to QQLinear layers for activation quantization. Only supported for models
            quantized with 'nvfp4' or 'mxfp8' modes. Default: ``False``.

    Returns:
        Tuple[nn.Module, TokenizerWrapper]: A tuple containing the loaded model and tokenizer.

    Raises:
        FileNotFoundError: If config file or safetensors are not found.
        ValueError: If model class or args class are not found.
    """
    force_download = kwargs.get("force_download", False)
    model_path = get_model_path(
        path_or_hf_repo, force_download=force_download, revision=revision
    )
    model = load_model(model_path, lazy, **kwargs)
    if adapter_path is not None:
        model = apply_lora_layers(model, adapter_path)
        model.eval()

    image_processor = load_image_processor(model_path, **kwargs)

    # Get the eos_token_id from the model config
    eos_token_id = getattr(model.config, "eos_token_id", None)

    processor = load_processor(model_path, True, eos_token_ids=eos_token_id, **kwargs)

    if image_processor is not None:
        processor.image_processor = image_processor

    return model, processor


def load_config(model_path: Union[str, Path], **kwargs) -> dict:
    """Load model configuration from a path or Hugging Face repo.

    Args:
        model_path: Local path or Hugging Face repo ID to load config from
        **kwargs: Additional keyword arguments to pass to the config loader

    Returns:
        dict: Model configuration

    Raises:
        FileNotFoundError: If config.json is not found at the path
    """
    if isinstance(model_path, str):
        model_path = get_model_path(model_path)

    try:
        with open(model_path / "config.json", encoding="utf-8") as f:
            config = json.load(f)

        generation_config_file = model_path / "generation_config.json"
        if generation_config_file.exists():
            generation_config = {}
            try:
                with open(generation_config_file, "r") as f:
                    generation_config = json.load(f)
            except json.JSONDecodeError:
                pass

            if eos_token_id := generation_config.get("eos_token_id", False):
                config["eos_token_id"] = eos_token_id

        return config

    except FileNotFoundError as exc:
        raise FileNotFoundError(f"Config not found at {model_path}") from exc


def load_image_processor(model_path: Union[str, Path], **kwargs) -> BaseImageProcessor:
    if isinstance(model_path, str):
        model_path = get_model_path(model_path)

    if not kwargs:
        config = load_config(model_path, trust_remote_code=True)
    else:
        config = load_config(model_path, **kwargs)

    model_class, _ = get_model_and_args(config)
    image_processor = None

    if hasattr(model_class, "ImageProcessor"):
        init_signature = inspect.signature(model_class.ImageProcessor.__init__)

        if "config" in init_signature.parameters:
            image_processor = model_class.ImageProcessor(config=config)
        else:
            image_processor = model_class.ImageProcessor()

    return image_processor


def load_processor(
    model_path, add_detokenizer=True, eos_token_ids=None, **kwargs
) -> ProcessorMixin:

    processor = AutoProcessor.from_pretrained(model_path, use_fast=True, **kwargs)
    if add_detokenizer:
        detokenizer_class = load_tokenizer(model_path, return_tokenizer=False)

        # Get the tokenizer object
        tokenizer_obj = (
            processor.tokenizer if hasattr(processor, "tokenizer") else processor
        )

        # Instantiate the detokenizer
        processor.detokenizer = detokenizer_class(tokenizer_obj)

        # Determine the EOS token IDs, prioritizing the function argument
        final_eos_token_ids = (
            eos_token_ids if eos_token_ids is not None else tokenizer_obj.eos_token_ids
        )

        # Create and assign the StoppingCriteria
        criteria = StoppingCriteria(final_eos_token_ids, tokenizer_obj)
        if hasattr(processor, "tokenizer"):
            processor.tokenizer.stopping_criteria = criteria
        else:
            processor.stopping_criteria = criteria

    return processor


def fetch_from_hub(
    model_path: Path, lazy: bool = False, **kwargs
) -> Tuple[nn.Module, dict, ProcessorMixin]:
    model = load_model(model_path, lazy, **kwargs)
    config = load_config(model_path, **kwargs)
    processor = load_processor(
        model_path,
        add_detokenizer=False,
        eos_token_ids=config.get("eos_token_id", None),
        **kwargs,
    )
    return model, config, processor


def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list:
    """
    Splits the weights into smaller shards.

    Args:
        weights (dict): Model weights.
        max_file_size_gb (int): Maximum size of each shard in gigabytes.

    Returns:
        list: List of weight shards.
    """
    max_file_size_bytes = max_file_size_gb << 30
    shards = []
    shard, shard_size = {}, 0
    for k, v in weights.items():
        if shard_size + v.nbytes > max_file_size_bytes:
            shards.append(shard)
            shard, shard_size = {}, 0
        shard[k] = v
        shard_size += v.nbytes
    shards.append(shard)
    return shards


def create_model_card(
    path: Union[str, Path], hf_path: Optional[Union[str, Path]] = None
):
    """
    Create model card for a converted MLX model.

    Args:
        path (Union[str, Path]): Local path to the converted model.
        hf_path (Optional[Union[str, Path]]): Original Hugging Face repo id or local path used for conversion.
    """
    from huggingface_hub import ModelCard, ModelCardData

    if hf_path is None:
        card = ModelCard.from_template(ModelCardData(language="en"))
    else:
        card = ModelCard.load(hf_path)
    card.data.library_name = "mlx"
    if card.data.pipeline_tag is None:
        card.data.pipeline_tag = "image-text-to-text"
    if card.data.tags is None:
        card.data.tags = ["mlx"]
    elif "mlx" not in card.data.tags:
        card.data.tags += ["mlx"]
    if hf_path is not None:
        card.data.base_model = str(hf_path)
    card.text = ""
    card.save(Path(path) / "README.md")


def upload_to_hub(path: str, upload_repo: str):
    """
    Uploads the model to Hugging Face hub.

    Args:
        path (str): Local path to the model.
        upload_repo (str): Name of the HF repo to upload to.
    """
    from huggingface_hub import HfApi, ModelCard, logging

    from . import __version__

    logging.set_verbosity_info()
    card_path = Path(path) / "README.md"
    card = ModelCard.load(card_path)

    hf_path = card.data.base_model

    if hf_path is not None:
        provenance = f"""
        This model was converted to MLX format from [`{hf_path}`](https://huggingface.co/{hf_path})
        using mlx-vlm version **{__version__}**.
        Refer to the [original model card](https://huggingface.co/{hf_path}) for more details on the model.
        """
    else:
        provenance = ""

    card.text = dedent(
        f"""
        # {upload_repo}
        {provenance}
        ## Use with mlx

        ```bash
        pip install -U mlx-vlm
        ```

        ```bash
        python -m mlx_vlm.generate --model {upload_repo} --max-tokens 100 --temperature 0.0 --prompt "Describe this image." --image <path_to_image>
        ```
        """
    )
    card.save(card_path)

    api = HfApi()
    api.create_repo(repo_id=upload_repo, exist_ok=True)
    api.upload_folder(
        folder_path=path,
        repo_id=upload_repo,
        repo_type="model",
    )
    print(f"Upload successful, go to https://huggingface.co/{upload_repo} for details.")


def apply_repetition_penalty(logits: mx.array, generated_tokens: Any, penalty: float):
    """
    Apply repetition penalty to specific logits based on the given context.

    Paper: https://arxiv.org/abs/1909.05858

    Args:
        logits (mx.array): The logits produced by the language model.
        generated_tokens (any): A list of N previous tokens.
        penalty (float): The repetition penalty factor to be applied.

    Returns:
        logits (mx.array): Logits with repetition penalty applied to generated tokens.
    """
    if len(generated_tokens) > 0:
        indices = mx.array([token for token in generated_tokens])
        selected_logits = logits[:, indices]
        selected_logits = mx.where(
            selected_logits < 0, selected_logits * penalty, selected_logits / penalty
        )
        logits[:, indices] = selected_logits
    return logits


def save_weights(
    save_path: Union[str, Path],
    model: nn.Module,
    *,
    donate_weights: bool = False,
) -> None:
    """Save model weights into specified directory."""
    if isinstance(save_path, str):
        save_path = Path(save_path)

    weights = dict(tree_flatten(model.parameters()))

    save_path.mkdir(parents=True, exist_ok=True)

    shards = make_shards(weights)
    shards_count = len(shards)
    shard_file_format = (
        "model-{:05d}-of-{:05d}.safetensors"
        if shards_count > 1
        else "model.safetensors"
    )

    total_size = sum(v.nbytes for v in weights.values())
    index_data = {"metadata": {"total_size": total_size}, "weight_map": {}}

    # Write the weights and make sure no references are kept other than the
    # necessary ones
    if donate_weights:
        model.update(tree_map(lambda _: mx.array([]), model.parameters()))

    weights.clear()
    del weights

    for i in range(len(shards)):
        shard = shards[i]
        shards[i] = None
        shard_name = shard_file_format.format(i + 1, shards_count)
        shard_path = save_path / shard_name

        mx.save_safetensors(str(shard_path), shard, metadata={"format": "mlx"})

        for weight_name in shard.keys():
            index_data["weight_map"][weight_name] = shard_name
        del shard

    index_data["weight_map"] = {
        k: index_data["weight_map"][k] for k in sorted(index_data["weight_map"])
    }

    with open(save_path / "model.safetensors.index.json", "w") as f:
        json.dump(
            index_data,
            f,
            indent=4,
        )


def save_config(
    config: dict,
    config_path: Union[str, Path],
) -> None:
    """Save the model configuration to the ``config_path``.

    The final configuration will be sorted before saving for better readability.

    Args:
        config (dict): The model configuration.
        config_path (Union[str, Path]): Model configuration file path.
    """
    # Clean unused keys
    config.pop("_name_or_path", None)
    config.pop("torch_dtype", None)

    # sort the config for better readability
    config = dict(sorted(config.items()))

    # write the updated config to the config_path (if provided)
    with open(config_path, "w") as fid:
        json.dump(config, fid, indent=4)


def load_image(image_source: Union[str, Path, BytesIO], timeout: int = 10):
    """
    Helper function to load an image from either a URL, file path, data URI,
    or BytesIO object.
    """
    import base64

    try:
        if not isinstance(image_source, (str, Path, BytesIO)):
            raise ValueError(
                f"Unsupported image source type: {type(image_source).__name__}"
            )
        if isinstance(image_source, str) and image_source.startswith("data:image/"):
            if "," not in image_source:
                raise ValueError("Invalid data URI format - missing comma separator")
            _, data = image_source.split(",", 1)
            image_source = BytesIO(base64.b64decode(data))
        if isinstance(image_source, str) and image_source.startswith(
            ("http://", "https://")
        ):
            response = requests.get(image_source, stream=True, timeout=timeout)
            response.raise_for_status()
            image_source = response.raw

        image = Image.open(image_source)
    except ValueError:
        raise
    except Exception as e:
        raise ValueError(f"Failed to load image from {image_source}: {e}") from e

    image = ImageOps.exif_transpose(image)
    return image.convert("RGB")


def resize_image(img, max_size):

    ratio = min(max_size[0] / img.width, max_size[1] / img.height)
    new_size = (int(img.width * ratio), int(img.height * ratio))
    return img.resize(new_size)


def process_image(img, resize_shape, image_processor):
    if isinstance(img, str):
        img = load_image(img)
    if hasattr(img, "mode") and img.mode != "RGB":
        img = img.convert("RGB")
    if resize_shape is not None and not isinstance(image_processor, BaseImageProcessor):
        img = resize_image(img, resize_shape)
    return img


def resample_audio(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    """Resample audio using linear interpolation."""
    if orig_sr == target_sr:
        return audio

    # Calculate the resampling ratio
    ratio = target_sr / orig_sr

    # Handle different audio shapes
    if audio.ndim == 1:
        # Mono audio - simple case
        new_length = int(len(audio) * ratio)
        old_indices = np.arange(len(audio))
        new_indices = np.linspace(0, len(audio) - 1, new_length)
        resampled = np.interp(new_indices, old_indices, audio)

    elif audio.ndim == 2:
        # Multi-channel audio - transpose to (samples, channels) if needed
        if audio.shape[0] < audio.shape[1]:
            audio = audio.T

        # Resample each channel
        n_samples, n_channels = audio.shape
        new_length = int(n_samples * ratio)
        old_indices = np.arange(n_samples)
        new_indices = np.linspace(0, n_samples - 1, new_length)

        resampled = np.zeros((new_length, n_channels))
        for i in range(n_channels):
            resampled[:, i] = np.interp(new_indices, old_indices, audio[:, i])
    else:
        raise ValueError(f"Audio array has unsupported shape: {audio.shape}")

    return resampled


def read_audio(file) -> tuple:
    """Read an audio file using miniaudio (or ffmpeg for m4a/aac/ogg/opus).

    Returns (samples_float32, sample_rate) where samples is always 2D (samples, channels).
    """
    import io as _io

    if isinstance(file, bytes):
        file = _io.BytesIO(file)

    # Check if ffmpeg is needed for certain formats
    use_ffmpeg = False
    if isinstance(file, (str, Path)):
        ext = Path(file).suffix.lstrip(".").lower()
        if ext in ("m4a", "aac", "ogg", "opus"):
            use_ffmpeg = True
    elif isinstance(file, _io.BytesIO):
        pos = file.tell()
        header = file.read(12)
        file.seek(pos)
        if header[4:8] == b"ftyp" or header[:4] == b"OggS":
            use_ffmpeg = True

    if use_ffmpeg:
        import json as _json
        import shutil
        import subprocess

        ffmpeg_path = shutil.which("ffmpeg")
        ffprobe_path = shutil.which("ffprobe")
        if ffmpeg_path is None:
            raise RuntimeError(
                "ffmpeg not found. Install it: brew install ffmpeg (macOS) "
                "or sudo apt install ffmpeg (Linux)"
            )

        if isinstance(file, _io.BytesIO):
            file.seek(0)
            input_data = file.read()
        else:
            input_data = None

        # Get info via ffprobe
        if ffprobe_path and input_data is not None:
            probe = subprocess.run(
                [
                    ffprobe_path,
                    "-v",
                    "quiet",
                    "-print_format",
                    "json",
                    "-show_streams",
                    "-select_streams",
                    "a:0",
                    "-i",
                    "pipe:0",
                ],
                input=input_data,
                capture_output=True,
            )
        elif ffprobe_path:
            probe = subprocess.run(
                [
                    ffprobe_path,
                    "-v",
                    "quiet",
                    "-print_format",
                    "json",
                    "-show_streams",
                    "-select_streams",
                    "a:0",
                    str(file),
                ],
                capture_output=True,
            )
        else:
            probe = None

        sample_rate, nchannels = 44100, 1
        if probe and probe.returncode == 0:
            info = _json.loads(probe.stdout)
            if info.get("streams"):
                stream = info["streams"][0]
                sample_rate = int(stream.get("sample_rate", 44100))
                nchannels = int(stream.get("channels", 1))

        # Decode via ffmpeg to raw PCM s16le
        cmd = [ffmpeg_path, "-v", "quiet"]
        if input_data is not None:
            cmd += ["-i", "pipe:0"]
        else:
            cmd += ["-i", str(file)]
        cmd += ["-f", "s16le", "-acodec", "pcm_s16le", "-ac", str(nchannels), "pipe:1"]

        result = subprocess.run(cmd, input=input_data, capture_output=True)
        if result.returncode != 0:
            raise RuntimeError(f"ffmpeg decoding failed: {result.stderr.decode()}")

        samples = np.frombuffer(result.stdout, dtype=np.int16)
    else:
        import miniaudio

        if isinstance(file, (str, Path)):
            info = miniaudio.get_file_info(str(file))
            decoded = miniaudio.decode_file(
                str(file),
                nchannels=info.nchannels,
                sample_rate=info.sample_rate,
            )
        elif isinstance(file, _io.BytesIO):
            file.seek(0)
            data = file.read()
            # Detect format from magic bytes
            if data[:4] == b"RIFF":
                info = miniaudio.wav_get_info(data)
            elif data[:3] == b"ID3" or data[:2] in (b"\xff\xfb", b"\xff\xfa"):
                info = miniaudio.mp3_get_info(data)
            elif data[:4] == b"fLaC":
                info = miniaudio.flac_get_info(data)
            else:
                info = miniaudio.vorbis_get_info(data)
            decoded = miniaudio.decode(
                data,
                nchannels=info.nchannels,
                sample_rate=info.sample_rate,
            )
        else:
            raise TypeError(f"Unsupported file type: {type(file)}")

        sample_rate = decoded.sample_rate
        nchannels = decoded.nchannels
        samples = np.array(decoded.samples, dtype=np.int16)

    # Reshape multi-channel and convert to float32
    if nchannels > 1:
        samples = samples.reshape(-1, nchannels)
    audio = samples.astype(np.float32) / 32768.0

    # Ensure always 2D
    if audio.ndim == 1:
        audio = audio[:, np.newaxis]

    return audio, sample_rate


def load_audio(
    file,
    sr: int,
    timeout: int = 10,
):
    """
    Helper function to load audio from either a URL, file path, or numpy array.
    """
    if isinstance(file, np.ndarray):
        return file
    if isinstance(file, Path):
        file = str(file)
    if isinstance(file, str) and file.startswith(("http://", "https://")):
        try:
            response = requests.get(file, stream=True, timeout=timeout)
            response.raise_for_status()
            audio, sample_rate = read_audio(BytesIO(response.content))
        except Exception as e:
            raise ValueError(
                f"Failed to load audio from URL: {file} with error {e}"
            ) from e
    else:
        audio, sample_rate = read_audio(file)

    if sample_rate != sr:
        audio = resample_audio(audio, sample_rate, sr)
    return np.array(audio).mean(axis=1)


def normalize_audio_features(features: mx.array) -> mx.array:
    """Normalize mel spectrogram features for lossy audio formats (e.g., MP3)."""
    return (features - mx.mean(features)) / (mx.std(features) + 1e-6)


def process_inputs(
    processor,
    prompts,
    images=None,
    audio=None,
    add_special_tokens=False,
    padding=True,
    padding_side="left",
    return_tensors="mlx",
    **kwargs,
):
    # Get the process method from the processor
    process_method = getattr(processor, "process", processor)
    parameters = inspect.signature(process_method).parameters

    # Prepare arguments
    args = {
        "text": prompts,
        "images": images,
        "padding": padding,
        "return_tensors": return_tensors,
    }
    if "padding_side" in parameters:
        args["padding_side"] = padding_side

    # Add special tokens if supported
    if "add_special_tokens" in parameters:
        args["add_special_tokens"] = add_special_tokens

    for param in parameters.keys():
        if param in kwargs.keys():
            args[param] = kwargs.get(param, None)

    # Add audio if provided and supported
    if audio is not None and len(audio) > 0:
        if "audio" in parameters:
            args["audio"] = audio
        elif "audios" in parameters:
            args["audios"] = audio
        else:
            raise ValueError(
                f"Processor {processor.__class__.__name__} does not support audio parameter"
            )

    return process_method(**args)


def process_inputs_with_fallback(
    processor,
    prompts,
    images,
    audio,
    add_special_tokens=False,
    return_tensors="mlx",
    **kwargs,
):
    # First attempt with specified return_tensors
    try:
        return process_inputs(
            processor,
            prompts=prompts,
            images=images,
            audio=audio,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            **kwargs,
        )
    except Exception as e:
        raise ValueError(f"Failed to process inputs with error: {e}")


def prepare_inputs(
    processor,
    images=None,
    audio=None,
    prompts=None,
    image_token_index=None,
    resize_shape=None,
    add_special_tokens=False,
    padding=True,
    padding_side="left",
    pad_to_uniform_size=False,
    return_tensors="mlx",
    **kwargs,
):

    has_images = images is not None and (
        not hasattr(images, "__len__") or len(images) > 0
    )
    has_audio = audio is not None and (not hasattr(audio, "__len__") or len(audio) > 0)
    if not has_images and not has_audio:
        tokenizer = (
            processor.tokenizer if hasattr(processor, "tokenizer") else processor
        )
        # Ensure pad_token exists when padding text-only inputs
        if padding and tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        inputs = tokenizer(
            prompts,
            add_special_tokens=add_special_tokens,
            padding=padding,
            padding_side=padding_side,
            return_tensors=return_tensors,
        )
        input_ids = (
            inputs.input_ids
            if isinstance(inputs.input_ids, mx.array)
            else mx.array(inputs.input_ids)
        )
        mask = (
            inputs.attention_mask
            if isinstance(inputs.attention_mask, mx.array)
            else mx.array(inputs.attention_mask)
        )
        return {
            "input_ids": input_ids,
            "attention_mask": mask,
        }

    # Process images
    if images is not None:
        if not isinstance(images, list):
            images = [images]

        image_processor = (
            processor.image_processor if hasattr(processor, "image_processor") else None
        )
        images = [process_image(img, resize_shape, image_processor) for img in images]

        # For batching, we need uniform image sizes. Instead of padding to the
        # largest image (which adds white borders that hurt accuracy), we resize
        # all images to the model's expected input size.
        if len(images) > 1 and pad_to_uniform_size:
            # Get target size from image processor if available
            target_size = None
            if image_processor is not None and hasattr(image_processor, "size"):
                size = image_processor.size
                if isinstance(size, tuple):
                    target_size = size
                elif isinstance(size, dict):
                    target_size = (size.get("height", 384), size.get("width", 384))
                elif isinstance(size, int):
                    target_size = (size, size)

            if target_size is not None:
                # Resize all images to the target size
                resized_images = []
                for img in images:
                    if img.size != (
                        target_size[1],
                        target_size[0],
                    ):  # PIL uses (width, height)
                        img = img.resize(
                            (target_size[1], target_size[0]), Image.Resampling.BICUBIC
                        )
                    resized_images.append(img)
                images = resized_images
            else:
                # Fallback: pad to largest size (original behavior)
                max_width = max(img.width for img in images)
                max_height = max(img.height for img in images)

                padded_images = []
                for img in images:
                    if img.width != max_width or img.height != max_height:
                        padded_img = Image.new(
                            "RGB", (max_width, max_height), (255, 255, 255)
                        )
                        x_offset = (max_width - img.width) // 2
                        y_offset = (max_height - img.height) // 2
                        padded_img.paste(img, (x_offset, y_offset))
                        padded_images.append(padded_img)
                    else:
                        padded_images.append(img)
                images = padded_images

    # Process audio
    audio_inputs = None
    audio_feature_lengths = None
    is_qwen3_omni_moe = False
    processor_class_name = (
        processor.__class__.__name__ if hasattr(processor, "__class__") else ""
    )
    if (
        "qwen3" in processor_class_name.lower()
        and "omni" in processor_class_name.lower()
    ):
        is_qwen3_omni_moe = True

    is_lossy_audio = False
    if audio is not None and len(audio) > 0:
        if not isinstance(audio, list):
            audio = [audio]

        # Check if any audio file is a lossy format (MP3, AAC, OGG, etc.)
        lossy_extensions = {".mp3", ".m4a"}
        is_lossy_audio = any(
            str(f).lower().endswith(tuple(lossy_extensions)) for f in audio
        )

        if len(audio) > 1:
            print(
                "\033[33mWarning\033[0m: Single prompt with multiple audio files is not supported yet. Using the first audio file.\n"
            )
            audio = audio[:1]

        if is_qwen3_omni_moe:
            audio_arrays = [
                load_audio(audio_file, sr=processor.feature_extractor.sampling_rate)
                for audio_file in audio
            ]
            audio_arrays = [
                audio_array.astype(np.float32) for audio_array in audio_arrays
            ]

            feature_extractor = getattr(processor, "feature_extractor", None)
            if feature_extractor is None:
                raise ValueError("Processor missing feature_extractor for audio prep.")

            audio_inputs = feature_extractor(
                audio_arrays,
                sampling_rate=feature_extractor.sampling_rate,
                padding=True,
                return_attention_mask=True,
            )

            audio_feature_lengths = np.sum(
                audio_inputs["attention_mask"], axis=-1, dtype=np.int32
            )
        else:
            feature_extractor = getattr(processor, "feature_extractor", None)
            sr = (
                getattr(feature_extractor, "sampling_rate", 16000)
                if feature_extractor is not None
                else 16000
            )
            audio = [load_audio(audio_file, sr=sr) for audio_file in audio]

    model_inputs = {}

    if hasattr(processor, "image_processor") and isinstance(
        processor.image_processor, BaseImageProcessor
    ):
        if not isinstance(prompts, list):
            prompts = [prompts]

        if processor.pad_token is None:
            processor.pad_token = processor.eos_token
        text_chunks = [
            [processor(chunk).input_ids for chunk in prompt.split("<image>")]
            for prompt in prompts
        ]

        # Find the maximum length for padding
        max_length = max(
            sum(len(chunk) for chunk in chunks) + 1 for chunks in text_chunks
        )

        # Pad and create input_ids
        input_ids = []
        for chunks in text_chunks:
            ids = chunks[0] + [image_token_index] + chunks[1]
            padding = [processor.pad_token_id] * (max_length - len(ids))
            input_ids.append(mx.array(ids + padding))

        model_inputs["input_ids"] = mx.array(input_ids)
        pixel_values = processor.image_processor.preprocess(images=images)
        model_inputs["pixel_values"] = mx.array(np.stack(pixel_values))
        model_inputs["attention_mask"] = mx.array(
            [(ids != processor.pad_token_id) for ids in input_ids]
        ).astype(mx.int32)

    else:
        if hasattr(processor, "tokenizer") and processor.tokenizer.pad_token is None:
            processor.tokenizer.pad_token = processor.tokenizer.eos_token

        inputs = process_inputs_with_fallback(
            processor,
            images=images,
            audio=audio,
            prompts=prompts,
            add_special_tokens=add_special_tokens,
            **kwargs,
        )

        if "images" in inputs:
            inputs["pixel_values"] = inputs["images"]
            inputs.pop("images")

        model_inputs["attention_mask"] = (
            mx.array(inputs["attention_mask"]) if "attention_mask" in inputs else None
        )

        # Convert inputs to model_inputs with mx.array if present
        for key, value in inputs.items():
            if key not in model_inputs:
                if value is None:
                    model_inputs[key] = value
                elif isinstance(value, (str, list, mx.array)):
                    model_inputs[key] = value
                else:
                    model_inputs[key] = mx.array(value)

    if audio_inputs is not None:
        model_inputs["input_features"] = mx.array(audio_inputs["input_features"])
        model_inputs["feature_attention_mask"] = mx.array(
            audio_inputs["attention_mask"]
        ).astype(mx.int32)
        model_inputs["audio_feature_lengths"] = mx.array(
            audio_feature_lengths, dtype=mx.int32
        )

    if is_lossy_audio and "input_features" in model_inputs:
        f = model_inputs["input_features"]
        if isinstance(f, list):
            model_inputs["input_features"] = [
                normalize_audio_features(mx.array(x)) for x in f
            ]
        else:
            model_inputs["input_features"] = normalize_audio_features(f)

    return model_inputs


def group_images_by_shape(
    images: List[Image.Image],
    disable_grouping: bool = False,
) -> Tuple[Dict[Tuple[int, int], List[Image.Image]], Dict[Tuple[int, int], List[int]]]:
    """
    Group images by their dimensions for efficient batch processing.

    Images with the same dimensions can be stacked and processed together,
    which is much faster than processing individually (especially on GPU).

    Args:
        images: List of PIL images to group
        disable_grouping: If True, each image gets its own group (useful for debugging)

    Returns:
        grouped_images: Dict mapping shape -> list of images with that shape
        grouped_indices: Dict mapping shape -> list of original indices

    Example:
        >>> images = [img_400x300, img_800x600, img_400x300_2]
        >>> grouped, indices = group_images_by_shape(images)
        >>> grouped
        {(300, 400): [img_400x300, img_400x300_2], (600, 800): [img_800x600]}
        >>> indices
        {(300, 400): [0, 2], (600, 800): [1]}
    """
    if disable_grouping:
        # Each image in its own group
        grouped_images = {}
        grouped_indices = {}
        for i, img in enumerate(images):
            shape = (img.height, img.width)
            # Make each shape unique by adding index
            unique_shape = (img.height, img.width, i)
            grouped_images[unique_shape] = [img]
            grouped_indices[unique_shape] = [i]
        return grouped_images, grouped_indices

    grouped_images: Dict[Tuple[int, int], List[Image.Image]] = {}
    grouped_indices: Dict[Tuple[int, int], List[int]] = {}

    for i, img in enumerate(images):
        shape = (img.height, img.width)
        if shape not in grouped_images:
            grouped_images[shape] = []
            grouped_indices[shape] = []
        grouped_images[shape].append(img)
        grouped_indices[shape].append(i)

    return grouped_images, grouped_indices


class StoppingCriteria:
    def __init__(self, eos_token_ids: List[int], tokenizer=None):

        if isinstance(eos_token_ids, int):
            self.eos_token_ids = [eos_token_ids]
        else:
            self.eos_token_ids = eos_token_ids

        self.tokenizer = tokenizer

    def add_eos_token_ids(self, new_eos_token_ids: Union[int, List[int]] = None):
        """
        Add new token IDs to the list of EOS token IDs.

        Args:
            new_eos_token_ids: Integer, string, or list of integers/strings representing token IDs to add.
                               If strings are provided, they will be converted to integers if possible.
        """
        if new_eos_token_ids is None:
            return

        if self.tokenizer is None:
            raise ValueError("Processor is not provided")

        if new_eos_token_ids is not None:
            if isinstance(new_eos_token_ids, str):
                new_eos_token_ids = [new_eos_token_ids]
            new_eos_token_ids = [
                self.tokenizer.encode(" " + token, add_special_tokens=False)[-1]
                for token in new_eos_token_ids
            ]
            self.eos_token_ids.extend(new_eos_token_ids)

    def reset(self, eos_token_ids: List[int] = None):
        eos_token_ids = (
            eos_token_ids if eos_token_ids is not None else self.tokenizer.eos_token_ids
        )

        if isinstance(eos_token_ids, int):
            eos_token_ids = [eos_token_ids]

        if self.eos_token_ids != eos_token_ids:
            self.eos_token_ids = eos_token_ids

    def __call__(self, input_ids: mx.array) -> bool:
        return input_ids in self.eos_token_ids


class ThinkingBudgetCriteria:
    """
    Enforces a budget on thinking tokens.

    Tracks tokens within thinking blocks (between start and end tokens) and
    forces a closing sequence (e.g. ``\\n</think>``) when budget is exceeded.
    """

    def __init__(
        self,
        tokenizer,
        thinking_budget: int,
        thinking_end_token: str = "</think>",
        thinking_start_token: Optional[str] = None,
        enable_thinking: bool = False,
    ):
        self.tokenizer = tokenizer
        self.thinking_budget = thinking_budget
        self.enable_thinking = enable_thinking

        # Resolve token IDs from strings
        self.thinking_end_token_id = tokenizer.encode(
            thinking_end_token, add_special_tokens=False
        )[-1]

        self.thinking_start_token_id = tokenizer.encode(
            thinking_start_token, add_special_tokens=False
        )[-1]

        self._forced_sequence: List[int] = []
        newline_ids = tokenizer.encode("\n", add_special_tokens=False)
        if newline_ids:
            self._forced_sequence.append(newline_ids[-1])
        self._forced_sequence.append(self.thinking_end_token_id)
        self._forced_index = 0

        self.in_thinking = self.enable_thinking
        self.thinking_token_count = 0
        self.budget_exceeded = False

    def reset_thinking_state(self):
        """Reset thinking state between generations."""
        self.in_thinking = self.enable_thinking
        self.thinking_token_count = 0
        self.budget_exceeded = False
        self._forced_index = 0

    def __call__(self, token_id: int) -> Optional[int]:
        """Process a token and return a forced token ID if budget exceeded, else None."""
        if self.enable_thinking and token_id == self.thinking_start_token_id:
            self.in_thinking = True
            return None

        if token_id == self.thinking_end_token_id:
            self.in_thinking = False
            self.budget_exceeded = False
            self._forced_index = 0
            return None

        if self.in_thinking:
            self.thinking_token_count += 1
            if self.thinking_token_count > self.thinking_budget:
                self.budget_exceeded = True

        if self.budget_exceeded and self._forced_index < len(self._forced_sequence):
            forced = self._forced_sequence[self._forced_index]
            self._forced_index += 1
            self.forced_token_id = forced
            return forced

        self.forced_token_id = None
        return None

    def apply_forced_token(self, next_y: mx.array) -> Optional[mx.array]:
        if self.forced_token_id is not None and self.enable_thinking:
            next_y = mx.array([self.forced_token_id])
            self.forced_token_id = None
            return next_y
        return next_y


def print_array_report(t: mx.array, label: Optional[str]) -> dict:
    """
    Return a dictionary report of an MLX array similar to PyTorch's tensor representation.
    Args:
        arr: MLX array to analyze
    Returns:
        Dictionary containing shape, dtype, value representation, and statistics
    """

    # Get basic statistics
    mean_val = mx.mean(t)
    std_val = mx.std(t)
    min_val = mx.min(t)
    max_val = mx.max(t)

    report = {
        "shape": f"{tuple(t.shape)}",
        "dtype": str(t.dtype),
        "value": repr(t),
        "mean": f"array({mean_val}, dtype={t.dtype})",
        "std": f"array({std_val}, dtype={t.dtype})",
        "min": f"array({min_val}, dtype={t.dtype})",
        "max": f"array({max_val}, dtype={t.dtype})",
        "label": label if label else "array",
    }

    # Print each field, handling 'value' specially
    print("{")
    for key, value in report.items():
        if key == "value":
            print(f" '{key}': {value},")  # No quotes around value
        else:
            print(f" '{key}': {repr(value)},")
    print("}")
    return report
