"""cider.ops — Low-level primitive API for W8A8 / W4A8 / per-group linear.

These functions return lazy mx.array nodes. Computation happens when
you call mx.eval() — fully compatible with MLX's graph-based execution.
"""

import re
import subprocess
from pathlib import Path
from typing import Optional

import mlx.core as mx
import numpy as np

# ── Kernel directory (shipped with the package) ─────────────────
_KERNEL_DIR: Optional[str] = None


def kernel_dir() -> str:
    """Return the absolute path to the bundled Metal kernels."""
    global _KERNEL_DIR
    if _KERNEL_DIR is None:
        _KERNEL_DIR = str(Path(__file__).parent / "kernels")
    return _KERNEL_DIR


# ── Extension loader ────────────────────────────────────────────
_ext = None


def _load_ext():
    global _ext
    if _ext is not None:
        return _ext
    import sys
    lib_dir = str(Path(__file__).parent / "lib")
    if lib_dir not in sys.path:
        sys.path.insert(0, lib_dir)
    try:
        import _cider_prim
        _ext = _cider_prim
        return _ext
    except ImportError:
        raise RuntimeError(
            "Cider C++ extension not available. INT8 TensorOps require Apple M5+. "
            "On M4 and below, use standard MLX inference instead."
        )


# ── Hardware detection ──────────────────────────────────────────

def is_available() -> bool:
    """Check if INT8 TensorOps are available (Apple M5+, Metal 4)."""
    try:
        chip = subprocess.run(
            ["sysctl", "-n", "machdep.cpu.brand_string"],
            capture_output=True, text=True, timeout=5
        ).stdout.strip()
    except Exception:
        return False
    m = re.match(r"Apple M(\d+)", chip)
    if not m or int(m.group(1)) < 5:
        return False
    try:
        _load_ext()
        return True
    except (ImportError, RuntimeError, Exception):
        return False


# ── Weight quantization helpers ─────────────────────────────────

def quantize_weight_int8(
    w: np.ndarray,
    clip_percentile: float = None,
) -> tuple:
    """Quantize FP16/FP32 weights to per-row symmetric INT8.

    Args:
        w: Weight matrix [N, K] as numpy array (N=out_features, K=in_features).
        clip_percentile: If set (e.g. 99.9), use per-row percentile instead of
            absmax to compute scale. Clips outliers before quantization, improving
            precision for layers with extreme weight outliers. Default None (absmax).

    Returns:
        (w_int8, scale_w) where w_int8 is [N, K] int8 and
        scale_w is [N] float32 (one scale per output channel).
    """
    w = w.astype(np.float32)
    if clip_percentile is not None:
        row_clip = np.percentile(np.abs(w), clip_percentile, axis=1)  # [N]
        row_clip = np.where(row_clip == 0, np.max(np.abs(w), axis=1), row_clip)
        scale = row_clip / 127.0
        scale = np.where(scale == 0, 1.0, scale)
        w = np.clip(w, -row_clip[:, np.newaxis], row_clip[:, np.newaxis])
    else:
        row_max = np.max(np.abs(w), axis=1)  # [N]
        scale = row_max / 127.0
        scale = np.where(scale == 0, 1.0, scale)
    w_int8 = np.clip(np.round(w / scale[:, np.newaxis]), -128, 127).astype(np.int8)
    return w_int8, scale.astype(np.float32)


def pack_weight_int4(
    w: np.ndarray,
    zero_point: int = 8,
) -> tuple:
    """Quantize FP16/FP32 weights to packed INT4 (symmetric, per-column).

    Args:
        w: Weight matrix [K, N] as numpy array. K must be even.
        zero_point: INT4 zero point (default 8 for signed range [-8, 7]).

    Returns:
        (packed_w, scale_w) where packed_w is [K//2, N] uint8
        (high nibble = even k, low nibble = odd k) and
        scale_w is [N] float32.
    """
    K, N = w.shape
    assert K % 2 == 0, f"K must be even, got {K}"
    w = w.astype(np.float32)
    col_max = np.max(np.abs(w), axis=0)
    scale = col_max / 7.0
    scale = np.where(scale == 0, 1.0, scale)
    w_q = np.clip(np.round(w / scale[np.newaxis, :]) + zero_point, 0, 15).astype(np.uint8)
    packed = (w_q[0::2, :] << 4) | w_q[1::2, :]
    return packed, scale.astype(np.float32)


# ── Primitive API ───────────────────────────────────────────────

def perchannel_linear(
    x: mx.array,
    w: mx.array,
    scale_w: mx.array,
    bias: Optional[mx.array] = None,
    stream: Optional[mx.Stream] = None,
) -> mx.array:
    """W8A8 per-channel quantized linear: y = dequant(quant_a(x) @ w_int8) + bias.

    Args:
        x: Input activations [M, K] float16 or bfloat16.
        w: INT8 weights [N, K] int8 (per-row quantized).
        scale_w: Per-row weight scales [N] float32.
        stream: Optional MLX stream.

    Returns:
        Output [M, N] matching input dtype.
    """
    ext = _load_ext()
    out_dtype = x.dtype
    kw = {"stream": stream} if stream is not None else {}
    N = w.shape[0]
    if bias is None:
        bias = mx.zeros((N,), dtype=mx.float16)
    result = ext.perchannel_linear(x, w, scale_w, bias, kernel_dir(), **kw)
    if out_dtype != mx.float16:
        result = result.astype(out_dtype, **kw)
    return result


# Shared placeholder for new_bias (V5 kernel ignores it; Metal needs valid buffer)
_shared_new_bias_cache = {}

def _get_shared_new_bias_placeholder(N: int, num_groups: int):
    key = (N, num_groups)
    if key not in _shared_new_bias_cache:
        _shared_new_bias_cache[key] = mx.zeros((N, num_groups), dtype=mx.float32)
    return _shared_new_bias_cache[key]


def pergroup_linear(
    x: mx.array,
    w: mx.array,
    scale_w: mx.array,
    group_size: int,
    bias: Optional[mx.array] = None,
    new_bias: Optional[mx.array] = None,
    stream: Optional[mx.Stream] = None,
) -> mx.array:
    """
    mlx native quantize format asymmetric affine
    quantize: q = clip(round((w - biases) / scales), 0, 2^b - 1), b = bits
    dequantize: w = q*scale + bias
    """
    """Per-group INT8 linear with optional bias.

    Dispatches internally:
      M > 1 → per-group GEMM (activation quantize + INT8 TensorOps)
      M == 1 → per-group MV (FP activation, weight dequant on-the-fly)

    Args:
        x: Input activations [M, K] float16 or bfloat16.
        w: INT8 weights [N, K] int8 (per-group symmetric quantized).
        scale_w: Per-group weight scales [num_groups, N] float32 (physically transposed for coalesced GPU access).
        group_size: Group size (64, 128, or 256).
        bias: Optional bias [N] float16. Default zeros.
        stream: Optional MLX stream.

    Returns:
        Output [M, N] matching input dtype.
    """
    ext = _load_ext()
    N = w.shape[0]
    num_groups = scale_w.shape[0] if scale_w.ndim == 2 else 1  # scale_w is [num_groups, N]
    if bias is None:
        bias = mx.zeros((N,), dtype=mx.float16)
    if new_bias is None:
        # V5 kernel ignores new_bias (symmetric quantization), but Metal
        # requires a valid buffer binding.  Use a tiny shared placeholder
        # instead of allocating (N, num_groups) every forward call.
        new_bias = _get_shared_new_bias_placeholder(N, num_groups)
    out_dtype = x.dtype
    kw = {"stream": stream} if stream is not None else {}
    # scale_w layout: [num_groups, N] physically contiguous. Kernel indexes as scale_w[g * N + n].

    result = ext.pergroup_linear(x, w, scale_w, bias, new_bias, group_size, kernel_dir(), **kw)
    if out_dtype != mx.float16:
        result = result.astype(out_dtype, **kw)
    return result


def w4a8_linear(
    x: mx.array,
    packed_w: mx.array,
    scale_w: mx.array,
    stream: Optional[mx.Stream] = None,
) -> mx.array:
    """W4A8 quantized linear: y = dequant(quant_a(x) @ unpack4(w)).

    Args:
        x: Input activations [M, K] float16.
        packed_w: Packed INT4 weights [K//2, N] uint8.
        scale_w: Per-column weight scales [N] float32.
        stream: Optional MLX stream.

    Returns:
        Output [M, N] float16.
    """
    ext = _load_ext()
    out_dtype = x.dtype
    kw = {"stream": stream} if stream is not None else {}
    result = ext.w4a8_linear(x, packed_w, scale_w, kernel_dir(), **kw)
    if out_dtype != mx.float16:
        result = result.astype(out_dtype, **kw)
    return result


def int8_matmul_int32(
    a: mx.array,
    b: mx.array,
    stream=None,
) -> mx.array:
    """Raw INT8×INT8→INT32 matmul (bit-exact, no dequant).

    Args:
        a: INT8 matrix [M, K].
        b: INT8 matrix [N, K] (transposed weight layout).
        stream: Optional MLX stream.

    Returns:
        INT32 result [M, N].
    """
    ext = _load_ext()
    kw = {"stream": stream} if stream is not None else {}
    return ext.int8_matmul_int32(a, b, kernel_dir(), **kw)