"""Independent full-prefix NumPy oracle for the native, incremental decoder."""

from concurrent.futures import ThreadPoolExecutor
import gc
import importlib
import os
import threading
import time

import numpy as np
import pytest


# A direct CMake build may expose the extension as _native before installation.
# Normal package/CI builds exercise the public cism._native installation path.
native = importlib.import_module(os.environ.get("CISM_NATIVE_TEST_MODULE", "cism._native"))


def tiny_model(model_type="llama", tied=False, hidden=24, intermediate=37):
    rng = np.random.default_rng(312)
    config = dict(
        model_type=model_type,
        hidden_size=hidden,
        intermediate_size=intermediate,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=8,
        vocab_size=43,
        max_position_embeddings=80,
        rms_norm_eps=1e-5,
        rope_theta=10000.0,
        tie_word_embeddings=tied,
        hidden_act="silu",
    )
    weights = {}

    def matrix(name, shape):
        weights[name] = rng.normal(0, 0.14, shape).astype(np.float32)

    def norm(name, width):
        weights[name] = rng.uniform(0.75, 1.25, width).astype(np.float32)

    matrix("model.embed_tokens.weight", (43, hidden))
    if not tied:
        matrix("lm_head.weight", (43, hidden))
    norm("model.norm.weight", hidden)
    for i in range(2):
        prefix = f"model.layers.{i}."
        norm(prefix + "input_layernorm.weight", hidden)
        norm(prefix + "post_attention_layernorm.weight", hidden)
        for name, shape in (
            ("q_proj", (32, hidden)),
            ("k_proj", (16, hidden)),
            ("v_proj", (16, hidden)),
            ("o_proj", (hidden, 32)),
        ):
            matrix(prefix + f"self_attn.{name}.weight", shape)
        if model_type == "qwen3":
            norm(prefix + "self_attn.q_norm.weight", 8)
            norm(prefix + "self_attn.k_norm.weight", 8)
        for name, shape in (
            ("gate_proj", (intermediate, hidden)),
            ("up_proj", (intermediate, hidden)),
            ("down_proj", (hidden, intermediate)),
        ):
            matrix(prefix + f"mlp.{name}.weight", shape)
    return config, weights


E2M1_HALF = np.array([0, 1, 2, 3, 4, 6, 8, 12], dtype=np.float32)


def encode_e4m3_numpy(values):
    """Round-to-nearest-even positive floats to E4M3, mirroring the C++ encoder."""
    values = np.asarray(values, dtype=np.float32)
    bits = np.zeros(values.shape, dtype=np.int64)
    normal = values >= np.float32(2.0 ** -6)
    with np.errstate(divide="ignore"):
        exponent = np.floor(np.log2(np.where(normal, values, np.float32(1)))).astype(np.int64)
    mantissa = values / (2.0 ** exponent)
    quantized_m = np.rint((mantissa - np.float32(1)) * 8).astype(np.int64)  # half-even
    field = exponent + 7 + (quantized_m >= 8)
    quantized_m = np.where(quantized_m >= 8, 0, quantized_m)
    bits[normal] = field[normal] * 8 + quantized_m[normal]
    subnormal = np.rint(values * 512.0).astype(np.int64)  # value / 2^-9
    bits[~normal] = np.minimum(subnormal[~normal], 0x7E)
    bits = np.minimum(bits, 0x7E)
    decoded = np.where(bits >= 8, (1 + (bits & 7) / 8) * 2.0 ** ((bits >> 3) - 7), (bits & 15) * 2.0 ** -9)
    return decoded.astype(np.float32)


def reconstructed_weights(weights, precision):
    result = {}
    for name, original in weights.items():
        if precision == "fp32" or original.ndim == 1:
            result[name] = original.copy()
            continue
        if precision == "hybrid-fp4" and ".mlp." in name:
            rows, cols = original.shape
            pad = (-cols) % 16
            blocks = np.pad(original, ((0, 0), (0, pad))).reshape(rows, -1, 16)
            amax = np.abs(blocks).max(axis=-1)
            # No 2^-6 floor here — matches the native packer, whose E4M3
            # encoder handles small scales via subnormals.
            ideal = np.where(amax == 0, np.float32(0), amax / 6)
            scale = encode_e4m3_numpy(ideal)
            magnitude = np.abs(blocks) / scale[..., None]
            # Compare against true E2M1 values (half table times one half).
            index = np.argmin(np.abs(magnitude[..., None] - E2M1_HALF * np.float32(0.5)), axis=-1)
            nibbles = np.where(np.sign(blocks) < 0, index | 8, index).astype(np.int64)
            decoded = np.where(np.sign(blocks) < 0, -E2M1_HALF[index], E2M1_HALF[index]) \
                * np.float32(0.5) * scale[..., None]
            result[name] = decoded.reshape(rows, -1)[:, :cols]
            continue
        int4 = precision == "hybrid-int4" and ".mlp." in name
        if int4:
            decoded = np.empty_like(original)
            for start in range(0, original.shape[1], 32):
                values = original[:, start : start + 32]
                m = values.shape[1]
                for r in range(values.shape[0]):
                    row = values[r]
                    # First-max wins (matches the native strict-greater scan).
                    idx = int(np.argmax(np.abs(row)))
                    extreme = np.float32(row[idx])
                    if extreme == 0:
                        d = np.float32(1.0)
                    else:
                        d = np.float32(extreme / np.float32(-8.0))
                    for t in range(m):
                        x = float(np.float32(row[t]) / d + np.float32(8.5))
                        xi = int(x)
                        if xi < 0:
                            xi = 0
                        elif xi > 15:
                            xi = 15
                        decoded[r, start + t] = np.float32(
                            (np.float32(xi) - np.float32(8.0)) * d)
            result[name] = decoded
            continue
        block = original.shape[1]
        maximum = np.float32(127)
        decoded = np.empty_like(original)
        for start in range(0, original.shape[1], block):
            values = original[:, start : start + block]
            largest = np.max(np.abs(values), axis=1, keepdims=True)
            scale = np.where(largest == 0, np.float32(1), largest / maximum)
            scale = np.maximum(scale, np.nextafter(np.float32(0), np.float32(1)))
            divided = values / scale
            fraction, integral = np.modf(np.abs(divided))
            rounded = np.copysign(integral + (fraction >= np.float32(0.5)), divided)
            quantized = np.clip(rounded, -maximum, maximum)
            decoded[:, start : start + block] = quantized * scale
        result[name] = decoded
    return result


def reference_logits(config, weights, tokens):
    """Vectorized dense attention over every prefix position, without a KV cache."""
    heads = config["num_attention_heads"]
    kv_heads = config["num_key_value_heads"]
    dim = config["head_dim"]
    length = len(tokens)

    def rms(values, weight):
        variance = np.mean(values.astype(np.float64) ** 2, axis=-1, keepdims=True)
        scale = (1 / np.sqrt(variance + config["rms_norm_eps"])).astype(np.float32)
        return (values * scale) * weight

    powers = np.arange(0, dim, 2, dtype=np.float32) / np.float32(dim)
    frequencies = np.float32(1) / (np.float32(config["rope_theta"]) ** powers)
    angles = np.arange(length, dtype=np.float32)[:, None] * frequencies[None, :]
    cosine, sine = np.cos(angles)[:, None, :], np.sin(angles)[:, None, :]

    def rotate(values):
        first, second = np.split(values, 2, axis=-1)
        return np.concatenate((first * cosine - second * sine, second * cosine + first * sine), axis=-1)

    x = weights["model.embed_tokens.weight"][tokens].copy()
    for layer in range(config["num_hidden_layers"]):
        p = f"model.layers.{layer}."
        normalized = rms(x, weights[p + "input_layernorm.weight"])
        q = (normalized @ weights[p + "self_attn.q_proj.weight"].T).reshape(length, heads, dim)
        k = (normalized @ weights[p + "self_attn.k_proj.weight"].T).reshape(length, kv_heads, dim)
        v = (normalized @ weights[p + "self_attn.v_proj.weight"].T).reshape(length, kv_heads, dim)
        if config["model_type"] == "qwen3":
            q = rms(q, weights[p + "self_attn.q_norm.weight"])
            k = rms(k, weights[p + "self_attn.k_norm.weight"])
        q = rotate(q)
        k = np.repeat(rotate(k), heads // kv_heads, axis=1)
        v = np.repeat(v, heads // kv_heads, axis=1)
        scores = np.einsum("thd,shd->hts", q, k) / np.sqrt(np.float32(dim))
        scores = np.where(np.triu(np.ones((length, length), dtype=bool), 1)[None], -np.inf, scores)
        probabilities = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
        probabilities /= np.sum(probabilities, axis=-1, keepdims=True)
        attention = np.einsum("hts,shd->thd", probabilities, v).reshape(length, heads * dim)
        x += attention @ weights[p + "self_attn.o_proj.weight"].T
        normalized = rms(x, weights[p + "post_attention_layernorm.weight"])
        gate = normalized @ weights[p + "mlp.gate_proj.weight"].T
        up = normalized @ weights[p + "mlp.up_proj.weight"].T
        silu = gate / (np.float32(1) + np.exp(-gate))
        x += (silu * up) @ weights[p + "mlp.down_proj.weight"].T
    final = rms(x, weights["model.norm.weight"])
    head = weights["model.embed_tokens.weight"] if config["tie_word_embeddings"] else weights["lm_head.weight"]
    return final[-1] @ head.T


@pytest.mark.parametrize("architecture", ["llama", "qwen3"])
@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
@pytest.mark.parametrize("tied", [False, True])
def test_logits_and_cached_generation(architecture, precision, tied):
    config, weights = tiny_model(architecture, tied)
    model = native.Model(config, weights, precision)
    reconstructed = reconstructed_weights(weights, precision)
    prompt = [2, 7, 1, 19, 4]
    for prefix in (prompt[:1], prompt[:3], prompt):
        actual = model.logits(prefix)
        assert actual.dtype == np.float32
        np.testing.assert_allclose(actual, reference_logits(config, reconstructed, prefix), atol=8e-6, rtol=8e-5)
    session = model.create_session(prompt, max_new_tokens=9)
    assert session.finish_reason == ""
    assert session.generated_tokens == 0
    assert session.next_tokens(0) == []
    expected = []
    for _ in range(9):
        logits = reference_logits(config, reconstructed, prompt + expected)
        expected.append(int(np.argmax(logits)))
    actual = session.next_tokens(2) + session.next_tokens(3) + session.next_tokens(100)
    assert actual == expected
    assert session.generated_tokens == 9
    assert session.finish_reason == "length"
    assert session.next_tokens(2) == []
    for index, token in enumerate(actual):
        assert token == int(np.argmax(model.logits(prompt + actual[:index])))


@pytest.mark.parametrize("architecture", ["llama", "qwen3"])
@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
def test_nll_matches_reference(architecture, precision):
    config, weights = tiny_model(architecture)
    model = native.Model(config, weights, precision)
    reconstructed = reconstructed_weights(weights, precision)
    tokens = [2, 7, 1, 19, 4, 11]
    # Row i of the reference (one prefix forward per position) predicts token i+1.
    rows = np.stack([reference_logits(config, reconstructed, tokens[: i + 1])
                     for i in range(len(tokens) - 1)])
    maximum = rows.max(axis=-1)
    log_total = np.log(np.exp(rows - maximum[:, None]).sum(axis=-1))
    targets = np.array(tokens[1:])
    expected = float(np.mean(maximum + log_total - rows[np.arange(len(tokens) - 1), targets]))
    actual = model.nll(tokens)
    assert isinstance(actual, float)
    assert actual == model.nll(tokens)  # deterministic
    assert abs(actual - expected) < 1e-4
    with pytest.raises(ValueError, match="at least 2 tokens"):
        model.nll(tokens[:1])


def expected_bytes(weights, precision):
    total = 0
    for name, array in weights.items():
        if array.ndim == 1 or precision == "fp32":
            total += array.nbytes
        elif precision == "hybrid-int4" and ".mlp." in name:
            rows, cols = array.shape
            blocks = (cols + 31) // 32
            total += rows * blocks * 16 + rows * blocks * 4
        elif precision == "hybrid-fp4" and ".mlp." in name:
            rows, cols = array.shape
            total += rows * ((cols + 1) // 2) + rows * ((cols + 15) // 16)
        else:
            total += array.size + array.shape[0] * 4
    return total


def test_fp4_element_lut_is_exact():
    expected = np.array([0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12])
    np.testing.assert_array_equal(np.asarray(native.fp4_element_lut()), expected)


@pytest.mark.parametrize("precision", ["int8", "hybrid-int4", "hybrid-fp4"])
def test_int8_activations_stay_close_and_deterministic(precision):
    # The pmaddwd quantized-activation path must engage (outputs move), stay
    # within a few percent of fp32 activations, and be bit-deterministic.
    # hidden=24/intermediate=37 exercise every partial-block tail.
    config, weights = tiny_model()
    prompt = [1, 3, 4, 2, 5]
    model = native.Model(config, weights, precision, 1)
    baseline = np.asarray(model.logits(prompt))
    model.set_act_precision("int8")
    quantized = np.asarray(model.logits(prompt))
    again = np.asarray(model.logits(prompt))
    assert not np.array_equal(baseline, quantized)
    np.testing.assert_allclose(quantized, baseline, rtol=0.03, atol=0.03)
    np.testing.assert_array_equal(quantized, again)
    # The constructor kwarg must land in the same place as the setter.
    via_ctor = native.Model(config, weights, precision, 1, "int8").logits(prompt)
    np.testing.assert_array_equal(quantized, np.asarray(via_ctor))


@pytest.mark.parametrize("precision", ["int8", "hybrid-int4", "hybrid-fp4"])
@pytest.mark.parametrize("act", ["fp32", "int8"])
def test_threaded_gemm_matches_single_thread(precision, act):
    # Wide enough rows*cols to engage the pooled path at threads=2; row
    # results are thread-independent, so threading must be bitwise exact.
    config, weights = tiny_model(hidden=256, intermediate=512)
    prompt = [1, 3, 4, 2, 5]
    single = np.asarray(native.Model(config, weights, precision, 1, act).logits(prompt))
    multi = np.asarray(native.Model(config, weights, precision, 2, act).logits(prompt))
    np.testing.assert_array_equal(single, multi)


def test_spec_stats_count_proposed_and_accepted():
    config, weights = tiny_model()
    model = native.Model(config, weights, "int8", 1)
    session = model.create_session([1, 2, 1, 2, 1, 2], max_new_tokens=16)
    assert session.spec_stats == (0, 0)
    first = session.next_tokens(16, 2)
    proposed, accepted = session.spec_stats
    assert proposed > 0
    assert 0 <= accepted <= proposed
    # Same model, fresh session: deterministic tally and tokens.
    again = model.create_session([1, 2, 1, 2, 1, 2], max_new_tokens=16)
    assert again.next_tokens(16, 2) == first
    assert again.spec_stats == (proposed, accepted)
    # Direct verification calls are tallied too, bounded by construction.
    direct = model.create_session([1, 2, 3], max_new_tokens=8)
    direct.verify([1, 2, 3])
    proposed2, accepted2 = direct.spec_stats
    assert proposed2 == 3
    assert 0 <= accepted2 <= 3


def test_engine_reports_spec_totals():
    from cism.engine import Engine

    class StubTokenizer:
        bos_token_id = 0
        eos_token_id = 1
        chat_template = None

        def encode(self, prompt, add_special_tokens=True):
            return list(prompt)

        def decode(self, ids, **kwargs):
            return ""

    config, weights = tiny_model()
    engine = Engine(native.Model(config, weights, "int8", 1), StubTokenizer(), config, "test")
    assert engine.generate([1, 2, 3], max_new_tokens=4, temperature=0).spec_proposed == 0
    chunks = list(engine.stream([1, 2, 1, 2, 1, 2], max_new_tokens=8, temperature=0, spec_k=2))
    assert sum(c.spec_proposed for c in chunks) > 0
    result = engine.generate([1, 2, 1, 2, 1, 2], max_new_tokens=8, temperature=0, spec_k=2)
    assert result.spec_proposed == sum(c.spec_proposed for c in chunks)
    assert result.spec_accepted == sum(c.spec_accepted for c in chunks)
    assert 0 <= result.spec_accepted <= result.spec_proposed


def test_fp4_e4m3_scale_encoding_matches_torch():
    torch = pytest.importorskip("torch")
    rng = np.random.default_rng(77)
    values = rng.uniform(2 ** -9, 400, 4096).astype(np.float32)
    encoded = np.array([native.fp4_encode_scale(float(value)) for value in values], dtype=np.uint8)
    # Native decode must agree with torch's view of the same bit pattern.
    torch_view = np.array([
        float(torch.from_numpy(np.array([bits], dtype=np.uint8)).view(torch.float8_e4m3fn))
        for bits in encoded], dtype=np.float32)
    native_view = np.array([native.fp4_decode_scale(int(bits)) for bits in encoded], dtype=np.float32)
    np.testing.assert_allclose(native_view, torch_view, rtol=0, atol=0)
    # And encoding a value must round-trip within E4M3 relative resolution.
    np.testing.assert_allclose(native_view, values, rtol=0.07)
    assert native.fp4_encode_scale(0.0) == 0
    assert native.fp4_encode_scale(1e30) == 0x7E
    assert native.fp4_decode_scale(0x7E) == 448.0


@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
def test_owned_weights_and_accounting(precision):
    config, weights = tiny_model("qwen3", tied=True)
    expected = expected_bytes(weights, precision)
    model = native.Model(config, weights, precision)
    before = model.logits([1, 2, 3])
    assert model.info["weight_bytes"] == expected
    assert model.info["precision"] == precision
    assert model.info["activation_dtype"] == "float32"
    assert model.info["act_precision"] == "fp32"
    assert model.info["threads"] == 1
    assert model.info["kernel"] in {"scalar", "avx2", "neon"}
    assert set(model.info["kernels"]) == {"fp32", "int8", "int4", "fp4", "variant"}
    assert all(name == model.info["kernel"] for key, name in model.info["kernels"].items()
               if key != "variant")
    assert model.info["kernels"]["variant"] in {"scalar", "avx2", "avx-vnni", "avx512-vnni", "neon"}
    model.set_act_precision("int8")
    assert model.info["activation_dtype"] == "int8-act"
    assert model.info["act_precision"] == "int8"
    model.set_act_precision("fp32")
    if expected_kernel := os.environ.get("CISM_NATIVE_EXPECT_KERNEL"):
        assert model.info["kernel"] == expected_kernel
    supplied_head = dict(weights, **{"lm_head.weight": weights["model.embed_tokens.weight"].copy()})
    assert native.Model(config, supplied_head, precision).info["weight_bytes"] == expected
    for array in weights.values():
        array.fill(0)
    config["hidden_size"] = 1
    np.testing.assert_array_equal(before, model.logits([1, 2, 3]))
    session = model.create_session([1, 2], max_new_tokens=3)
    del model
    gc.collect()
    assert len(session.next_tokens(9)) == 3


def test_sampling_chunks_seed_and_filters():
    config, weights = tiny_model()
    model = native.Model(config, weights)
    options = dict(max_new_tokens=20, temperature=0.9, top_p=0.8, top_k=11, seed=876)
    first = model.create_session([3, 1], **options).next_tokens(20)
    session = model.create_session([3, 1], **options)
    second = session.next_tokens(1) + session.next_tokens(7) + session.next_tokens(12)
    assert first == second
    other = model.create_session([3, 1], **dict(options, seed=877)).next_tokens(20)
    assert first != other
    greedy = model.create_session([3], max_new_tokens=10).next_tokens(10)
    assert model.create_session([3], max_new_tokens=10, temperature=1, top_k=1).next_tokens(10) == greedy
    assert model.create_session([3], max_new_tokens=10, temperature=1, top_p=1e-12).next_tokens(10) == greedy
    assert model.create_session([3], max_new_tokens=10, temperature=0, top_k=1, top_p=0.2).next_tokens(10) == greedy
    prefix = [3, 1]
    for token in first:
        logits = model.logits(prefix)
        order = np.argsort(-logits, kind="stable")[:11]
        probabilities = np.exp((logits[order].astype(np.float64) - logits[order[0]]) / 0.9)
        probabilities /= probabilities.sum()
        count = np.searchsorted(np.cumsum(probabilities), 0.8) + 1
        assert token in order[:count]
        prefix.append(token)


@pytest.mark.parametrize("architecture", ["llama", "qwen3"])
@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
@pytest.mark.parametrize("spec_k", [2, 4])
def test_speculative_greedy_matches_plain_decode(architecture, precision, spec_k):
    config, weights = tiny_model(architecture, tied=True)
    model = native.Model(config, weights, precision)
    # A repetitive prompt makes the n-gram drafter fire and multi-accept.
    prompt = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4]
    plain = model.create_session(prompt, max_new_tokens=24, eos_token_ids=[]).next_tokens(24)
    session = model.create_session(prompt, max_new_tokens=24, eos_token_ids=[])
    chunked = session.next_tokens(3, spec_k=spec_k) + session.next_tokens(21, spec_k=spec_k)
    assert chunked == plain
    assert session.finish_reason == "length"


def test_speculative_draft_hits_and_misses():
    config, weights = tiny_model("llama", tied=True)
    model = native.Model(config, weights)
    prompt = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4]

    def py_draft(history, k):
        """Independent reimplementation of the v4 prompt-lookup drafter:
        most-recent strictly-inside occurrence per position, iteratively
        extended (fixes v1 truncation at history end); stop when the
        context never occurred."""
        out = []
        size = len(history)
        if k == 0 or size < 2:
            return out
        for _ in range(k):
            total = size + len(out)
            placed = False
            for pattern in range(min(3, total - 1), 0, -1):
                if pattern >= size:
                    continue
                ctx = (history + out)[total - pattern:total]
                for s in range(size - pattern - 1, -1, -1):
                    if history[s:s + pattern] == ctx:
                        out.append(history[s + pattern])
                        placed = True
                        break
                if placed:
                    break
            if not placed:
                break
        return out

    session = model.create_session(prompt, max_new_tokens=8, eos_token_ids=[])
    history = list(prompt)
    for _ in range(3):
        generated = session.next_tokens(1)
        history.extend(generated)
        assert session.draft(3) == py_draft(history, 3)
    assert session.draft(0) == []
    empty = model.create_session([1], max_new_tokens=4, eos_token_ids=[])
    assert empty.draft(4) == []  # History too short to find a pattern.


def test_speculative_verify_accept_and_reject():
    config, weights = tiny_model("llama", tied=True)
    model = native.Model(config, weights)
    prompt = [1, 2, 3]
    plain = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[]).next_tokens(16)

    # Feeding the true continuation must fully accept (plus bonus) and land on
    # exactly the same greedy sequence as plain decoding.
    session = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[])
    out = []
    while len(out) < 16:
        take = min(3, 16 - len(out))
        out += session.verify(plain[len(out):len(out) + take])
    assert out == plain

    # Wrong candidates must be rejected and corrected, still matching greedy.
    session = model.create_session(prompt, max_new_tokens=16, eos_token_ids=[])
    out = []
    while len(out) < 16:
        take = min(3, 16 - len(out))
        wrong = [(token + 7) % config["vocab_size"] for token in plain[len(out):len(out) + take]]
        out += session.verify(wrong)
    assert out == plain


def test_speculative_rejection_path_and_boundaries():
    config, weights = tiny_model("qwen3", tied=True)
    model = native.Model(config, weights)
    # spec_k validation happens at the session call boundary.
    session = model.create_session([1, 2], max_new_tokens=10, eos_token_ids=[])
    with pytest.raises(ValueError, match="spec_k"):
        session.next_tokens(4, spec_k=1)
    with pytest.raises(ValueError, match="spec_k"):
        session.next_tokens(4, spec_k=17)
    # EOS inside an accepted speculative block must stop cleanly.
    eos = int(np.argmax(model.logits([1, 2])))
    reference = model.create_session([1, 2], max_new_tokens=12, eos_token_ids=[eos])
    expected = reference.next_tokens(12)
    assert reference.finish_reason in ("stop", "length")
    speculative = model.create_session([1, 2], max_new_tokens=12, eos_token_ids=[eos])
    actual = speculative.next_tokens(12, spec_k=3)
    assert actual == expected
    assert speculative.finish_reason == reference.finish_reason
    # Token budget boundary: speculation never overshoots max_new_tokens.
    exact = model.create_session([1, 2], max_new_tokens=2, eos_token_ids=[])
    tokens = exact.next_tokens(2, spec_k=4)
    assert len(tokens) == 2 and exact.generated_tokens == 2
    assert exact.finish_reason in ("stop", "length")


def test_eos_lengths_and_cancellation():
    config, weights = tiny_model()
    model = native.Model(config, weights)
    eos = int(np.argmax(model.logits([1, 2])))
    session = model.create_session([1, 2], max_new_tokens=1, eos_token_ids=[eos])
    assert session.next_tokens(20) == [eos]
    assert session.finish_reason == "stop"
    assert session.generated_tokens == 1
    session.cancel()
    assert session.finish_reason == "stop"
    empty = model.create_session([1], max_new_tokens=0)
    assert empty.finish_reason == "length"
    assert empty.next_tokens(1) == []
    cancelled = model.create_session([1], max_new_tokens=7)
    cancelled.cancel()
    assert cancelled.finish_reason == "cancelled"
    assert cancelled.next_tokens(7) == []
    assert cancelled.generated_tokens == 0
    # EOS in the prompt is ordinary context, not a generated stopping event.
    session = model.create_session([eos], max_new_tokens=2, eos_token_ids=[eos])
    assert session.finish_reason == ""
    assert len(session.next_tokens(1)) == 1
    config["max_position_embeddings"] = 3
    small = native.Model(config, weights)
    assert len(small.create_session([1, 2], max_new_tokens=1).next_tokens(2)) == 1
    assert small.logits([1, 2, 3]).shape == (43,)
    with pytest.raises(ValueError, match="exceeds"):
        small.create_session([1, 2], max_new_tokens=2)


@pytest.mark.parametrize("prompt", [[], [-1], [43], [2**62], [1.0], [True]])
def test_invalid_prompt(prompt):
    config, weights = tiny_model()
    model = native.Model(config, weights)
    with pytest.raises((ValueError, TypeError, OverflowError)):
        model.create_session(prompt, max_new_tokens=1)
    with pytest.raises((ValueError, TypeError, OverflowError)):
        model.logits(prompt)


@pytest.mark.parametrize("options", [
    {"max_new_tokens": -1}, {"max_new_tokens": 2**62},
    {"temperature": -1}, {"temperature": float("nan")}, {"temperature": float("inf")},
    {"top_p": 0}, {"top_p": -0.1}, {"top_p": 1.1}, {"top_p": float("nan")},
    {"top_k": -1}, {"top_k": 44}, {"eos_token_ids": [-1]}, {"eos_token_ids": [43]},
    {"seed": -1},
])
def test_invalid_sampling(options):
    config, weights = tiny_model()
    model = native.Model(config, weights)
    with pytest.raises((ValueError, TypeError, OverflowError)):
        model.create_session([1], **dict({"max_new_tokens": 1}, **options))
    with pytest.raises(ValueError):
        model.create_session([1], max_new_tokens=1).next_tokens(-1)


@pytest.mark.parametrize("change", [
    {"model_type": "qwen3_moe"}, {"hidden_size": 0}, {"hidden_size": 65537},
    {"head_dim": 7}, {"num_key_value_heads": 3}, {"num_attention_heads": -1},
    {"num_hidden_layers": 257}, {"intermediate_size": True}, {"vocab_size": 4.5},
    {"max_position_embeddings": 1048577}, {"rms_norm_eps": 0}, {"rms_norm_eps": float("nan")},
    {"rope_theta": float("inf")}, {"rope_theta": -1},
    {"attention_bias": True}, {"mlp_bias": True}, {"qkv_bias": True},
    {"bias": True}, {"is_encoder_decoder": True}, {"add_cross_attention": True},
    {"rope_interleaved": True}, {"num_experts": 8},
    {"quantization_config": {"quant_method": "gptq"}}, {"compression_config": {"format": "packed"}},
    {"rope_scaling": {"rope_type": "linear", "factor": 2}},
    {"rope_parameters": {"rope_type": "yarn"}}, {"rope_type": "llama3"},
    {"partial_rotary_factor": 0.5}, {"rotary_pct": 0.5},
    {"sliding_window": 64}, {"use_sliding_window": True},
    {"layer_types": ["full_attention", "sliding_attention"]},
    {"layer_types": ["full_attention"]},
    {"hidden_act": "gelu"}, {"pretraining_tp": 2}, {"attention_multiplier": 0.3},
    {"hidden_size": 65536, "vocab_size": 2000000},
])
def test_invalid_config(change):
    config, weights = tiny_model()
    config.update(change)
    with pytest.raises((ValueError, TypeError, OverflowError)):
        native.Model(config, weights)


@pytest.mark.parametrize("mutation", ["missing", "extra", "bias", "shape", "rank", "dtype", "noncontiguous", "nan", "inf", "list", "endian"])
def test_invalid_weights(mutation):
    config, weights = tiny_model()
    name = "model.embed_tokens.weight"
    if mutation == "missing":
        del weights[name]
    elif mutation == "extra":
        weights["model.unknown.weight"] = np.zeros(1, np.float32)
    elif mutation == "bias":
        weights["model.layers.0.self_attn.q_proj.bias"] = np.zeros(32, np.float32)
    elif mutation == "shape":
        weights[name] = weights[name][:-1].copy()
    elif mutation == "rank":
        weights[name] = weights[name].ravel()
    elif mutation == "dtype":
        weights[name] = weights[name].astype(np.float64)
    elif mutation == "noncontiguous":
        weights[name] = np.asfortranarray(weights[name])
    elif mutation == "nan":
        weights[name][0, 0] = np.nan
    elif mutation == "inf":
        weights[name][0, 0] = np.inf
    elif mutation == "list":
        weights[name] = weights[name].tolist()
    elif mutation == "endian":
        weights[name] = weights[name].astype(">f4")
    with pytest.raises(ValueError):
        native.Model(config, weights)


def test_tied_and_untied_head_validation():
    config, weights = tiny_model(tied=True)
    weights["lm_head.weight"] = weights["model.embed_tokens.weight"].copy()
    weights["lm_head.weight"][0, 0] += 1
    with pytest.raises(ValueError, match="tied"):
        native.Model(config, weights)
    del weights["lm_head.weight"]
    config["tie_word_embeddings"] = False
    with pytest.raises(ValueError, match="lm_head"):
        native.Model(config, weights)
    with pytest.raises(ValueError, match="precision"):
        native.Model(config, weights, "int4")


@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
def test_zero_weights_and_tie_breaking(precision):
    config, weights = tiny_model()
    for value in weights.values():
        value.fill(0)
    model = native.Model(config, weights, precision)
    np.testing.assert_array_equal(model.logits([1, 2]), np.zeros(43, np.float32))
    assert model.create_session([1], max_new_tokens=3).next_tokens(3) == [0, 0, 0]
    assert model.create_session([1], max_new_tokens=3, temperature=1, top_k=1).next_tokens(3) == [0, 0, 0]


def test_unaligned_contiguous_input():
    config, weights = tiny_model()
    name = "model.embed_tokens.weight"
    original = weights[name]
    buffer = bytearray(original.nbytes + 1)
    unaligned = np.ndarray(original.shape, dtype=np.float32, buffer=buffer, offset=1)
    unaligned[:] = original
    weights[name] = unaligned
    np.testing.assert_allclose(native.Model(config, weights).logits([1, 3]), reference_logits(config, weights, [1, 3]), atol=8e-6)


def test_concurrent_sessions_and_serialized_calls():
    config, weights = tiny_model("qwen3")
    model = native.Model(config, weights, "hybrid-int4")
    expected = model.create_session([1, 2], max_new_tokens=24).next_tokens(24)
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [executor.submit(lambda: model.create_session([1, 2], max_new_tokens=24).next_tokens(24)) for _ in range(4)]
        assert all(future.result(timeout=20) == expected for future in futures)
        shared = model.create_session([1, 2], max_new_tokens=24)
        futures = [executor.submit(shared.next_tokens, 6) for _ in range(4)]
        results = [future.result(timeout=20) for future in futures]
    assert sorted(map(tuple, results)) == sorted(tuple(expected[i : i + 6]) for i in range(0, 24, 6))
    assert shared.generated_tokens == 24
    assert shared.finish_reason == "length"


def test_gil_release_and_atomic_cancel():
    config, weights = tiny_model()
    config["max_position_embeddings"] = 100000
    model = native.Model(config, weights)
    session = model.create_session([1, 2], max_new_tokens=99990)
    started = threading.Event()

    def decode():
        started.set()
        return session.next_tokens(99990)

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(decode)
        assert started.wait(timeout=5)
        session.cancel()
        # This property must not wait on the mutex while holding the GIL.
        reason = session.finish_reason
        output = future.result(timeout=20)
    assert reason == "cancelled"
    assert session.generated_tokens == len(output)
    assert len(output) < 99990


def test_cancel_during_lazy_prefill():
    config, weights = tiny_model()
    config["max_position_embeddings"] = 100000
    model = native.Model(config, weights)
    session = model.create_session([1, 2] * 25000, max_new_tokens=2)
    started = threading.Event()

    def decode():
        started.set()
        return session.next_tokens(2)

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(decode)
        assert started.wait(timeout=5)
        # Allow native prefill to run; a 50k-token quadratic prefill cannot finish
        # in this interval. The monitor must be able to acquire the released GIL.
        time.sleep(0.02)
        assert not future.done()
        session.cancel()
        assert future.result(timeout=20) == []
    assert session.finish_reason == "cancelled"
    assert session.generated_tokens == 0
    assert session.next_tokens(2) == []


def test_lazy_prefill_owns_prompt():
    config, weights = tiny_model()
    model = native.Model(config, weights)
    prompt = [1, 2, 3]
    expected = int(np.argmax(model.logits(prompt)))
    session = model.create_session(prompt, max_new_tokens=1)
    prompt[:] = [-1]
    assert session.next_tokens(1) == [expected]


def test_kv_limit_rejected_before_allocation():
    config, weights = tiny_model()
    config["max_position_embeddings"] = 1048576
    config["num_hidden_layers"] = 256
    for layer in range(2, 256):
        for name, value in list(weights.items()):
            if name.startswith("model.layers.0."):
                weights[name.replace("model.layers.0.", f"model.layers.{layer}.")] = value
    model = native.Model(config, weights)
    with pytest.raises(ValueError, match="KV cache"):
        model.create_session([1], max_new_tokens=1048575)


@pytest.mark.parametrize("architecture", ["llama", "qwen3"])
@pytest.mark.parametrize("kv_heads", [1, 4])
def test_multi_query_and_multi_head_attention(architecture, kv_heads):
    config, weights = tiny_model(architecture)
    config["num_key_value_heads"] = kv_heads
    rng = np.random.default_rng(913)
    for layer in range(config["num_hidden_layers"]):
        for projection in ("k_proj", "v_proj"):
            weights[f"model.layers.{layer}.self_attn.{projection}.weight"] = rng.normal(
                0, 0.14, (kv_heads * config["head_dim"], config["hidden_size"])
            ).astype(np.float32)
    model = native.Model(config, weights)
    prompt = [1, 4, 6, 7]
    np.testing.assert_allclose(model.logits(prompt), reference_logits(config, weights, prompt), atol=8e-6, rtol=8e-5)
    expected = []
    for _ in range(5):
        expected.append(int(np.argmax(reference_logits(config, weights, prompt + expected))))
    assert model.create_session(prompt, max_new_tokens=5).next_tokens(5) == expected


def test_inference_failure_terminates_session():
    config, weights = tiny_model()
    model = native.Model(config, weights)
    # A finite FP32 input can still overflow a dot product. Refuse non-finite output.
    weights["model.layers.0.self_attn.q_proj.weight"].fill(np.finfo(np.float32).max)
    weights["model.layers.0.self_attn.k_proj.weight"].fill(np.finfo(np.float32).max)
    invalid = native.Model(config, weights)
    with pytest.raises(RuntimeError, match="non-finite"):
        invalid.logits([1, 2])
    # Finite weights pass construction; the first real inference is lazy.
    session = invalid.create_session([1, 2], max_new_tokens=2)
    assert session.finish_reason == ""
    assert session.next_tokens(0) == []
    with pytest.raises(RuntimeError, match="non-finite"):
        session.next_tokens(1)
    assert session.finish_reason == "cancelled"
    assert session.generated_tokens == 0
    assert session.next_tokens(1) == []
    cancelled = invalid.create_session([1, 2], max_new_tokens=2)
    cancelled.cancel()
    assert cancelled.next_tokens(1) == []
    assert cancelled.finish_reason == "cancelled"
    empty = invalid.create_session([1, 2], max_new_tokens=0)
    assert empty.next_tokens(1) == []
    assert empty.finish_reason == "length"
    assert np.isfinite(model.logits([1, 2])).all()


@pytest.mark.parametrize("architecture", ["llama", "qwen3"])
@pytest.mark.parametrize("precision", ["int8", "hybrid-int4"])
@pytest.mark.parametrize("width", [1, 7, 8, 9, 31, 32, 33, 63, 64, 65, 127, 256])
def test_quantized_vector_boundaries(architecture, precision, width):
    config, weights = tiny_model(architecture, hidden=width, intermediate=width + 2)
    reconstructed = reconstructed_weights(weights, precision)
    model = native.Model(config, weights, precision)
    prompt = [1, 5, 2]
    expected = []
    for _ in range(4):
        prefix = prompt + expected
        target = reference_logits(config, reconstructed, prefix)
        np.testing.assert_allclose(model.logits(prefix), target, rtol=0, atol=1e-4)
        expected.append(int(np.argmax(target)))
    session = model.create_session(prompt, max_new_tokens=4)
    assert session.next_tokens(1) + session.next_tokens(3) == expected


def test_int4_signed_nibbles_and_distinct_block_scales():
    config, weights = tiny_model("qwen3", hidden=65, intermediate=97)
    for name, weight in weights.items():
        if ".mlp." not in name:
            continue
        # Include every representable quantized value, distinct adjacent nibbles,
        # negative/positive row extrema, zero blocks, and multiple block scales.
        rows, cols = np.indices(weight.shape)
        weight[:] = (((cols + rows) % 15) - 7) * np.where((cols // 32) % 2, 0.007, 0.02)
        weight[0, :32] = 0
        weight[1, 32:64] = 0
    reconstructed = reconstructed_weights(weights, "hybrid-int4")
    model = native.Model(config, weights, "hybrid-int4")
    for prompt in ([1], [3, 9, 2], [2, 7, 5, 6, 1]):
        np.testing.assert_allclose(model.logits(prompt), reference_logits(config, reconstructed, prompt), rtol=0, atol=1e-4)


@pytest.mark.parametrize("threads", [1, 2, 4])
def test_multithreaded_matches_single_thread_bitwise(threads):
    config, weights = tiny_model("qwen3", tied=True)
    single = native.Model(config, weights, "hybrid-int4", threads=1)
    multi = native.Model(config, weights, "hybrid-int4", threads=threads)
    assert multi.info["threads"] == threads
    # Row slices execute identical per-row dot kernels; results are exact.
    for prompt in ([1], [3, 9, 2], [2, 7, 5, 6, 1]):
        np.testing.assert_array_equal(multi.logits(prompt), single.logits(prompt))
    session_single = single.create_session([1, 5, 2], max_new_tokens=8, temperature=0.9,
                                           top_p=0.9, top_k=20, seed=42, eos_token_ids=[])
    session_multi = multi.create_session([1, 5, 2], max_new_tokens=8, temperature=0.9,
                                         top_p=0.9, top_k=20, seed=42, eos_token_ids=[])
    np.testing.assert_array_equal(session_multi.next_tokens(8), session_single.next_tokens(8))


def test_invalid_thread_counts():
    config, weights = tiny_model("llama", tied=True)
    for threads in (0, 65, -1):
        with pytest.raises(ValueError, match="threads"):
            native.Model(config, weights, "fp32", threads=threads)
    # Floats are strictly rejected at the binding boundary (TypeError), while
    # booleans are integer subclasses: True is a valid single-threaded model.
    with pytest.raises(TypeError):
        native.Model(config, weights, "fp32", threads=2.5)
    assert native.Model(config, weights, "fp32", threads=True).info["threads"] == 1


@pytest.mark.parametrize("precision", ["fp32", "int8", "hybrid-int4", "hybrid-fp4"])
def test_page_locking_and_touch(precision):
    config, weights = tiny_model("qwen3", tied=True)
    model = native.Model(config, weights, precision)
    info_bytes = model.info["weight_bytes"]
    assert not model.info["pages_locked"]
    # Touch sums every stored weight byte, including quantized buffers, scales,
    # and the rotary table, which weight_bytes excludes (head_dim/2 floats).
    first = model.touch()
    rotary_bytes = (config["head_dim"] // 2) * 4
    assert info_bytes < first <= info_bytes + rotary_bytes
    assert model.touch() == first
    assert model.scan() == first
    assert model.scan() == model.scan()
    locked = model.lock_pages()
    if locked:
        assert locked <= info_bytes + rotary_bytes
        assert model.info["pages_locked"]
        with pytest.raises(ValueError, match="already locked"):
            model.lock_pages()
        # Page sharing makes exact byte equality impossible; small neighboring
        # regions occupy the same OS page, so unlock can only report <= locked.
        unlocked = model.unlock_pages()
        assert 0 < unlocked <= locked
    assert not model.info["pages_locked"]
    with pytest.raises(ValueError, match="not locked"):
        model.unlock_pages()
    # Inference still works after lock/unlock cycles on all storage formats.
    model.create_session([1], max_new_tokens=2, temperature=0, top_p=1, top_k=0,
                         seed=0, eos_token_ids=[]).next_tokens(2)