"""cider.convert — One-line model conversion for W8A8 prefill acceleration.

Usage:
    from cider import convert_model, set_mode

    model, proc = load("model_path")
    convert_model(model)         # Patch all Linear layers in-place

works with standard MLX LLM/VLM architectures — tested on Qwen, Llama, Qwen3-VL
Supports float16 and bfloat16 models automatically.
"""

import gc
import time

import mlx.nn as nn

from . import ops
from .nn import CiderLinear

# Re-export mode control from nn (single source of truth)


# ── Model conversion ───────────────────────────────────────────
_TARGET_TYPES = (nn.Linear, nn.QuantizedLinear)


def _convert_children(module, counter, clip_percentile=None):
    """Walk module.children(), replace Linear/QuantizedLinear in-place."""
    for name, child in module.children().items():
        if isinstance(child, _TARGET_TYPES):
            setattr(module, name, CiderLinear.from_float(child, clip_percentile=clip_percentile))
            counter[0] += 1
            if counter[0] % 28 == 0:
                gc.collect()
        elif isinstance(child, list):
            for i, item in enumerate(child):
                if isinstance(item, _TARGET_TYPES):
                    child[i] = CiderLinear.from_float(item, clip_percentile=clip_percentile)
                    counter[0] += 1
                    if counter[0] % 28 == 0:
                        gc.collect()
                elif isinstance(item, nn.Module):
                    _convert_children(item, counter, clip_percentile=clip_percentile)
        elif isinstance(child, nn.Module):
            _convert_children(child, counter, clip_percentile=clip_percentile)
        # Skip dict/other non-Module children (e.g. rope_scaling)


def convert_model(
    model: nn.Module,
    *,
    clip_percentile: float = None,
    verbose: bool = True,
) -> dict:
    """Convert all Linear/QuantizedLinear layers to CiderLinear (in-place).

    After conversion:
    - set_mode("prefill") → all linears use W8A8 INT8 TensorOps
    - set_mode("decode")  → all linears use original weights (no overhead)

    Args:
        model: Any MLX nn.Module (Qwen3-VL, Llama, Mistral, etc.).
        verbose: Print conversion summary.

    Returns:
        dict with stats: n_converted, elapsed_s.

    Example:
        from cider import convert_model, set_mode

        model, proc = load("model_path")
        stats = convert_model(model)

        set_mode("prefill")
        # ... run prefill (W8A8, ~15-19% faster) ...

        set_mode("decode")
        # ... run decode (original weights, optimal for single token) ...
    """
    if not ops.is_available():
        raise RuntimeError(
            "W8A8 INT8 TensorOps not available. Requires Apple M5+ with Metal 4."
        )

    t0 = time.perf_counter()
    counter = [0]
    _convert_children(model, counter, clip_percentile=clip_percentile)

    elapsed = time.perf_counter() - t0
    n = counter[0]
    stats = {"n_converted": n, "elapsed_s": elapsed}

    if verbose:
        print(
            f"[cider] Converted {n} layers to CiderLinear "
            f"in {elapsed:.1f}s"
        )

    return stats