"""MLX backend: convert models to MLX format with mixed-precision quantization."""

import json
import os
import shutil
from typing import Callable, Optional


def convert_llm_to_mlx(
    hf_path: str,
    output_path: str,
    quant_predicate: Optional[Callable] = None,
    q_bits: int = 4,
    q_group_size: int = 64,
    metadata: Optional[dict] = None,
) -> str:
    """Convert a HuggingFace LLM to quantized MLX format.

    Uses mlx-lm's convert() with optional custom quant_predicate for mixed-precision.
    """
    from mlx_lm import convert

    # mlx-lm refuses to overwrite existing dirs, so remove first
    if os.path.exists(output_path):
        shutil.rmtree(output_path)

    print(f"  Converting {hf_path} → {output_path}")
    print(f"  Quantization: {'mixed-precision (OptiQ)' if quant_predicate else f'uniform {q_bits}-bit'}")

    convert(
        hf_path=hf_path,
        mlx_path=output_path,
        quantize=True,
        q_bits=q_bits,
        q_group_size=q_group_size,
        quant_predicate=quant_predicate,
    )

    # Save OptiQ metadata alongside the model
    if metadata:
        meta_path = os.path.join(output_path, "optiq_metadata.json")
        with open(meta_path, "w") as f:
            json.dump(metadata, f, indent=2, default=str)
        print(f"  Saved optimization metadata to {meta_path}")

    # Report model size
    total_size = sum(
        os.path.getsize(os.path.join(output_path, f))
        for f in os.listdir(output_path)
        if f.endswith((".safetensors", ".npz"))
    )
    print(f"  Output model size: {total_size / 1024 / 1024:.1f} MB")

    return output_path


def convert_llm_uniform(
    hf_path: str,
    output_path: str,
    bits: int = 4,
    group_size: int = 64,
) -> str:
    """Convert LLM with uniform quantization (baseline)."""
    return convert_llm_to_mlx(
        hf_path=hf_path,
        output_path=output_path,
        quant_predicate=None,
        q_bits=bits,
        q_group_size=group_size,
    )


def convert_llm_static_mixed(
    hf_path: str,
    output_path: str,
    recipe: str = "mixed_3_6",
) -> str:
    """Convert LLM with mlx-lm's built-in static mixed-precision recipe."""
    from mlx_lm import convert

    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    print(f"  Converting {hf_path} → {output_path} (static recipe: {recipe})")

    convert(
        hf_path=hf_path,
        mlx_path=output_path,
        quantize=True,
        quant_predicate=recipe,
    )

    total_size = sum(
        os.path.getsize(os.path.join(output_path, f))
        for f in os.listdir(output_path)
        if f.endswith((".safetensors", ".npz"))
    )
    print(f"  Output model size: {total_size / 1024 / 1024:.1f} MB")
    return output_path


def _pt_to_mlx_vlm_path(pt_name: str) -> str:
    """Map PyTorch layer name to mlx-vlm layer path.

    PyTorch (HuggingFace):          MLX (mlx-vlm):
      model.visual.blocks.X.Y    → vision_tower.blocks.X.Y
      model.visual.merger.Y      → vision_tower.merger.Y
      model.language_model.Y     → language_model.model.Y
      lm_head                    → (tied, not quantized)
    """
    if pt_name == "lm_head":
        return "lm_head"
    if pt_name.startswith("model.visual."):
        return pt_name.replace("model.visual.", "vision_tower.", 1)
    if pt_name.startswith("model.language_model."):
        return pt_name.replace("model.language_model.", "language_model.model.", 1)
    return pt_name


def make_vlm_quant_predicate(
    configs: dict[str, int],
    default_bits: int = 4,
    group_size: int = 64,
) -> Callable:
    """Build a quant_predicate for mlx-vlm from OptiQ per-layer configs.

    Args:
        configs: Dict mapping PyTorch layer names to bit-widths.
        default_bits: Default bits for layers not in configs.
        group_size: Quantization group size.

    Returns:
        Callable suitable for mlx_vlm.convert(quant_predicate=...).
    """
    # Build MLX path → bits mapping
    mlx_configs = {}
    for pt_name, bits in configs.items():
        mlx_path = _pt_to_mlx_vlm_path(pt_name)
        mlx_configs[mlx_path] = bits

    def predicate(path: str, module) -> bool | dict:
        if path in mlx_configs:
            bits = mlx_configs[path]
            return {"bits": bits, "group_size": group_size}
        # For unmapped layers, use default
        return {"bits": default_bits, "group_size": group_size}

    return predicate


def convert_vlm_to_mlx(
    hf_path: str,
    output_path: str,
    quant_predicate: Optional[Callable] = None,
    q_bits: int = 4,
    q_group_size: int = 64,
    metadata: Optional[dict] = None,
) -> str:
    """Convert a HuggingFace VLM to quantized MLX format.

    Uses mlx-vlm's convert() with optional custom quant_predicate for mixed-precision.
    """
    from mlx_vlm import convert

    if os.path.exists(output_path):
        shutil.rmtree(output_path)

    print(f"  Converting VLM {hf_path} → {output_path}")
    print(f"  Quantization: {'mixed-precision (OptiQ)' if quant_predicate else f'uniform {q_bits}-bit'}")

    convert(
        hf_path=hf_path,
        mlx_path=output_path,
        quantize=True,
        q_bits=q_bits,
        q_group_size=q_group_size,
        quant_predicate=quant_predicate,
    )

    if metadata:
        meta_path = os.path.join(output_path, "optiq_metadata.json")
        with open(meta_path, "w") as f:
            json.dump(metadata, f, indent=2, default=str)
        print(f"  Saved optimization metadata to {meta_path}")

    total_size = sum(
        os.path.getsize(os.path.join(output_path, f))
        for f in os.listdir(output_path)
        if f.endswith((".safetensors", ".npz"))
    )
    print(f"  Output model size: {total_size / 1024 / 1024:.1f} MB")

    return output_path


def make_audio_quant_predicate(
    configs: dict[str, int],
    default_bits: int = 4,
    group_size: int = 64,
    skip_prefixes: list[str] | None = None,
) -> "Callable":
    """Build a class_predicate for mx.nn.quantize from OptiQ per-layer configs.

    Layer names are MLX-native (e.g. model.layers.0.mlp.gate_proj).
    By default, skips audio_tower layers (kept in BF16).

    Returns:
        Callable suitable for mlx.nn.quantize(class_predicate=...).
    """
    import mlx.nn as nn

    if skip_prefixes is None:
        skip_prefixes = ["audio_tower."]

    def predicate(path: str, module):
        if not isinstance(module, nn.Linear):
            return False
        # Skip layers that should stay in full precision
        for prefix in skip_prefixes:
            if path.startswith(prefix):
                return False
        if path in configs:
            bits = configs[path]
            return {"bits": bits, "group_size": group_size}
        return {"bits": default_bits, "group_size": group_size}

    return predicate


def convert_audio_to_mlx(
    model,
    output_path: str,
    class_predicate: "Optional[Callable]" = None,
    q_bits: int = 4,
    q_group_size: int = 64,
    metadata: Optional[dict] = None,
    source_config_path: str | None = None,
) -> str:
    """Quantize an MLX audio model (Qwen3-ASR) and save to disk.

    Takes an already-loaded MLX model, applies quantization,
    and saves weights + config.

    Args:
        model: MLX model (e.g. from mlx_audio.stt.utils.load_model)
        output_path: Directory to save quantized model
        class_predicate: Per-layer quantization predicate
        q_bits: Default bit-width for uniform quantization
        q_group_size: Quantization group size
        metadata: OptiQ metadata to save alongside model
        source_config_path: Path to source config.json to copy
    """
    import mlx.nn as nn
    import mlx.core as mx
    from mlx.utils import tree_flatten

    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    os.makedirs(output_path, exist_ok=True)

    print(f"  Converting audio model → {output_path}")
    print(f"  Quantization: {'mixed-precision (OptiQ)' if class_predicate else f'uniform {q_bits}-bit'}")

    # Apply quantization (skip audio_tower by default to match mlx-audio format)
    if class_predicate:
        nn.quantize(model, class_predicate=class_predicate)
    else:
        def uniform_pred(path, module):
            if not isinstance(module, nn.Linear):
                return False
            if path.startswith("audio_tower."):
                return False
            return {"bits": q_bits, "group_size": q_group_size}
        nn.quantize(model, class_predicate=uniform_pred)

    # Save weights (as model.safetensors to match mlx-audio's expected naming)
    weights = dict(tree_flatten(model.parameters()))
    mx.save_safetensors(os.path.join(output_path, "model.safetensors"), weights)

    # Collect per-layer quantization info from the quantized model
    per_layer_bits = {}
    if class_predicate:
        for path, module in model.named_modules():
            if hasattr(module, "bits"):
                per_layer_bits[path] = int(module.bits)

    # Copy config and tokenizer files from source
    if source_config_path and os.path.exists(source_config_path):
        src_dir = os.path.dirname(source_config_path)
        # Copy config with quantization info
        with open(source_config_path) as f:
            config = json.load(f)
        quant_config = {"bits": q_bits, "group_size": q_group_size}
        # Embed per-layer overrides so mlx-audio's apply_quantization()
        # creates QuantizedLinear with the correct bits for each layer
        for layer_path, bits in per_layer_bits.items():
            if bits != q_bits:
                quant_config[layer_path] = {"bits": bits, "group_size": q_group_size}
        config["quantization"] = quant_config
        with open(os.path.join(output_path, "config.json"), "w") as f:
            json.dump(config, f, indent=2)
        # Copy tokenizer and other necessary files (skip weight/index files)
        copy_exts = (".json", ".txt", ".model", ".tiktoken", ".py")
        skip_files = {"config.json", "model.safetensors"}
        for fname in os.listdir(src_dir):
            if fname in skip_files:
                continue
            # Skip safetensors index and shard files (we save our own weights)
            if "safetensors" in fname:
                continue
            if any(fname.endswith(ext) for ext in copy_exts):
                src_file = os.path.join(src_dir, fname)
                dst_file = os.path.join(output_path, fname)
                if os.path.isfile(src_file) and not os.path.exists(dst_file):
                    shutil.copy2(src_file, dst_file)

    # Save per-layer quantization map in metadata
    if class_predicate and per_layer_bits:
        if metadata is None:
            metadata = {}
        metadata["per_layer_bits"] = per_layer_bits

    # Save metadata
    if metadata:
        meta_path = os.path.join(output_path, "optiq_metadata.json")
        with open(meta_path, "w") as f:
            json.dump(metadata, f, indent=2, default=str)

    total_size = get_model_size_mb(output_path)
    print(f"  Output model size: {total_size:.1f} MB")

    return output_path


def get_model_size_mb(path: str) -> float:
    """Get total model file size in MB."""
    total = 0
    for f in os.listdir(path):
        fp = os.path.join(path, f)
        if os.path.isfile(fp) and f.endswith((".safetensors", ".npz", ".bin")):
            total += os.path.getsize(fp)
    return total / (1024 * 1024)
