"""Qwen3 forward in MLX with explicit per-layer attention masks.

Attention modules adapted from mlx-lm (MIT License, Copyright Apple Inc.),
https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/qwen3.py, reduced to the
inference paths this artifact needs: every call takes an explicit mask, so the stock
causal-mask construction is deliberately absent. ``Denoiser`` is the numpy-facing wrapper
the sampler talks to; masks cross the boundary in the blocked convention (``True`` =
blocked) and are inverted exactly once, here.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np

ARCHITECTURE_KEYS = (
    'vocab_size', 'hidden_size', 'num_hidden_layers', 'num_attention_heads',
    'num_key_value_heads', 'head_dim', 'intermediate_size', 'rms_norm_eps', 'rope_theta',
)


@dataclass
class ModelArgs:
    vocab_size: int
    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    head_dim: int
    intermediate_size: int
    rms_norm_eps: float
    rope_theta: float


class KVCache:
    """Append-only key/value buffer with logical truncation via ``trim``.

    ``offset`` is the logical length; buffers grow in 256-token steps and are overwritten
    past ``offset`` after a trim, which is what lets the denoising loop re-feed a changing
    block against a fixed prefix.
    """

    step = 256

    def __init__(self) -> None:
        self.keys: mx.array | None = None
        self.values: mx.array | None = None
        self.offset = 0

    def update_and_fetch(self, keys: mx.array, values: mx.array) -> tuple[mx.array, mx.array]:
        end = self.offset + keys.shape[2]
        if self.keys is None or end > self.keys.shape[2]:
            batch, n_kv_heads, _, head_dim = keys.shape
            size = ((end + self.step - 1) // self.step) * self.step
            grown_keys = mx.zeros((batch, n_kv_heads, size, head_dim), keys.dtype)
            grown_values = mx.zeros((batch, n_kv_heads, size, head_dim), values.dtype)
            if self.keys is not None:
                grown_keys[..., : self.offset, :] = self.keys[..., : self.offset, :]
                grown_values[..., : self.offset, :] = self.values[..., : self.offset, :]
            self.keys, self.values = grown_keys, grown_values
        self.keys[..., self.offset : end, :] = keys
        self.values[..., self.offset : end, :] = values
        self.offset = end
        return self.keys[..., :end, :], self.values[..., :end, :]

    def trim(self, count: int) -> int:
        count = min(self.offset, count)
        self.offset -= count
        return count


class Attention(nn.Module):
    def __init__(self, args: ModelArgs) -> None:
        super().__init__()
        self.n_heads = args.num_attention_heads
        self.n_kv_heads = args.num_key_value_heads
        self.scale = args.head_dim**-0.5
        dim = args.hidden_size
        self.q_proj = nn.Linear(dim, self.n_heads * args.head_dim, bias=False)
        self.k_proj = nn.Linear(dim, self.n_kv_heads * args.head_dim, bias=False)
        self.v_proj = nn.Linear(dim, self.n_kv_heads * args.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * args.head_dim, dim, bias=False)
        self.q_norm = nn.RMSNorm(args.head_dim, eps=args.rms_norm_eps)
        self.k_norm = nn.RMSNorm(args.head_dim, eps=args.rms_norm_eps)
        self.rope = nn.RoPE(args.head_dim, traditional=False, base=args.rope_theta)

    def __call__(self, x: mx.array, mask: mx.array, cache: KVCache | None = None) -> mx.array:
        batch, length, _ = x.shape
        queries = self.q_proj(x).reshape(batch, length, self.n_heads, -1)
        keys = self.k_proj(x).reshape(batch, length, self.n_kv_heads, -1)
        values = self.v_proj(x).reshape(batch, length, self.n_kv_heads, -1)
        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
        keys = self.k_norm(keys).transpose(0, 2, 1, 3)
        values = values.transpose(0, 2, 1, 3)
        offset = cache.offset if cache is not None else 0
        queries = self.rope(queries, offset=offset)
        keys = self.rope(keys, offset=offset)
        if cache is not None:
            keys, values = cache.update_and_fetch(keys, values)
        output = mx.fast.scaled_dot_product_attention(
            queries, keys, values, scale=self.scale, mask=mask
        )
        output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return self.o_proj(output)


class MLP(nn.Module):
    def __init__(self, args: ModelArgs) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))


class TransformerBlock(nn.Module):
    def __init__(self, args: ModelArgs) -> None:
        super().__init__()
        self.self_attn = Attention(args)
        self.mlp = MLP(args)
        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
        self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)

    def __call__(self, x: mx.array, mask: mx.array, cache: KVCache | None = None) -> mx.array:
        x = x + self.self_attn(self.input_layernorm(x), mask, cache)
        return x + self.mlp(self.post_attention_layernorm(x))


class Qwen3Model(nn.Module):
    def __init__(self, args: ModelArgs) -> None:
        super().__init__()
        self.args = args
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)

    def __call__(
        self, input_ids: mx.array, mask: mx.array, caches: list[KVCache] | None = None
    ) -> mx.array:
        """Hidden states under an explicit mask; ``mask`` broadcasts to ``[B, H, Q, K]``."""

        hidden = self.embed_tokens(input_ids)
        caches = caches or [None] * len(self.layers)
        for layer, cache in zip(self.layers, caches):
            hidden = layer(hidden, mask, cache)
        return self.norm(hidden)


class Denoiser:
    """Single-sequence forward with position gathering and forbidden-logit masking.

    Mirrors the torch ``Qwen3Denoiser`` contract on numpy tensors: token ids in, fp32
    logits out, gathered at the requested row positions before the tied head so the full
    vocabulary is never materialized for unneeded rows.
    """

    def __init__(self, model: Qwen3Model, config: dict) -> None:
        self.model = model
        self.config = config
        mdlm = config['mdlm']
        forbidden = np.zeros(config['vocab_size'], dtype=bool)
        forbidden[np.asarray(mdlm['forbidden_output_token_ids'])] = True
        forbidden[mdlm['forbidden_output_from'] :] = True
        self._forbidden = forbidden

    def new_caches(self) -> list[KVCache]:
        return [KVCache() for _ in range(len(self.model.layers))]

    @staticmethod
    def trim_caches(caches: list[KVCache], keep: int) -> None:
        excess = caches[0].offset - keep
        if excess > 0:
            for cache in caches:
                cache.trim(excess)

    def logits(
        self,
        token_ids,
        blocked: np.ndarray,
        positions,
        caches: list[KVCache] | None = None,
    ) -> np.ndarray | None:
        """Logits ``[len(positions), vocab]`` for the fed rows; ``None`` primes a cache only.

        ``blocked`` is ``[Q, K]`` with ``True`` marking a blocked key; with ``caches`` the
        key axis must cover ``offset + Q`` exactly, matching the cached forward contract of
        the torch backend.
        """

        ids = np.asarray(token_ids, dtype=np.int64)
        query_len = ids.shape[0]
        expected_keys = (caches[0].offset if caches else 0) + query_len
        if blocked.shape != (query_len, expected_keys):
            raise ValueError(
                f'blocked must be [{query_len}, {expected_keys}], got {blocked.shape}'
            )
        mask = mx.array(~blocked)[None, None, :, :]
        hidden = self.model(mx.array(ids[None, :]), mask, caches)
        if positions is None:
            mx.eval(hidden)
            return None
        taken = mx.take(hidden[0], mx.array(np.asarray(positions, dtype=np.int32)), axis=0)
        logits = self.model.embed_tokens.as_linear(taken).astype(mx.float32)
        out = np.array(logits)
        out[:, self._forbidden] = np.finfo(np.float32).min
        return out


def load_denoiser(directory: str | Path, *, dtype: str | None = None) -> tuple[Denoiser, dict]:
    """Build the model from ``config.json`` + ``model.safetensors`` in ``directory``.

    ``dtype`` recasts the stored weights on load (e.g. ``'float32'`` for parity checks);
    the shipped file is fp16.
    """

    directory = Path(directory)
    config = json.loads((directory / 'config.json').read_text())
    model = Qwen3Model(ModelArgs(**{key: config[key] for key in ARCHITECTURE_KEYS}))
    weights = mx.load(str(directory / 'model.safetensors'))
    if dtype is not None:
        weights = {name: array.astype(getattr(mx, dtype)) for name, array in weights.items()}
    model.load_weights(list(weights.items()))
    mx.eval(model.parameters())
    return Denoiser(model, config), config