"""Packed serving: compressed tensors stay compressed; base loaders must never run."""
import json

import numpy as np
import pytest
import torch

mx = pytest.importorskip("mlx.core")
nn = pytest.importorskip("mlx.nn")
pytest.importorskip("mlx_lm")

from mlx.utils import tree_flatten
from mlx_lm.models import qwen3_5
from tokenizers import Tokenizer, models, pre_tokenizers
from transformers import PreTrainedTokenizerFast

from kev.checkpoint import Checkpoint, LoadOptions
from kev.model import PointerHead, SPECIAL
from kev.packed_model import PackedLinear
from scripts.build_packed_model import build, pack_native


def test_native_codes_and_compensation():
    rng = np.random.default_rng(8)
    for bits in (2, 3, 4):
        codes = rng.integers(0, 3 if bits == 2 else 1 << bits, size=(8, 64), dtype=np.uint8)
        scale = mx.array(rng.uniform(.01, .1, (8, 1)).astype(np.float16))
        bias = -scale
        q = PackedLinear(64, 8, bits, 64, rank=2)
        q.weight, q.scales, q.biases = mx.array(pack_native(codes, bits)), scale, bias
        q.comp_a = mx.array(rng.normal(0, .01, (2, 64))).astype(mx.bfloat16)
        q.comp_b = mx.array(rng.normal(0, .01, (8, 2))).astype(mx.bfloat16)
        x = mx.array(rng.normal(size=(3, 64))).astype(mx.bfloat16)
        dense = mx.dequantize(q.weight, scale, bias, group_size=64, bits=bits)
        decoded = (mx.array(codes).astype(mx.float32) * scale.astype(mx.float32) + bias.astype(mx.float32)).astype(mx.float16)
        np.testing.assert_array_equal(np.asarray(dense), np.asarray(decoded))
        expected = (x.astype(mx.float32) @ dense.astype(mx.float32).T +
                    (x.astype(mx.float32) @ q.comp_a.astype(mx.float32).T) @ q.comp_b.astype(mx.float32).T).astype(mx.bfloat16)
        np.testing.assert_allclose(np.asarray(q(x).astype(mx.float32)), np.asarray(expected.astype(mx.float32)), atol=.008, rtol=.008)
        assert q.weight.dtype == mx.uint32 and q.weight.nbytes == codes.size * bits // 8


def tiny_bundle(tmp_path):
    source = tmp_path / "source"
    source.mkdir()
    config = {"model_type": "qwen3_5_text", "hidden_size": 64, "intermediate_size": 128, "num_hidden_layers": 2,
              "num_attention_heads": 2, "num_key_value_heads": 1, "head_dim": 32, "vocab_size": 32,
              "linear_num_value_heads": 2, "linear_num_key_heads": 2, "linear_key_head_dim": 128,
              "linear_value_head_dim": 128, "linear_conv_kernel_dim": 4, "full_attention_interval": 2,
              "tie_word_embeddings": True, "layer_types": ["linear_attention", "full_attention"]}
    (source / "config.json").write_text(json.dumps(config))
    vocab = {w: i for i, w in enumerate(["<unk>", "<pad>", *SPECIAL, "billing", "shipping", "charged", "twice"])}
    tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
    PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="<unk>", pad_token="<pad>", additional_special_tokens=SPECIAL).save_pretrained(source)
    head = PointerHead(64, 8)
    torch.save({"base": "offline/does-not-exist", "head": head.state_dict(), "head_dim": 8, "temperature": 2.4}, source / "head.pt")
    lm = qwen3_5.Model(qwen3_5.ModelArgs.from_dict(config))
    rng = np.random.default_rng(12)
    tensors, kinds, groups = {}, {}, {}
    modules = dict(lm.named_modules())
    for name, value in tree_flatten(lm.parameters()):
        module_name = name.removesuffix(".weight")
        if name.endswith(".weight") and isinstance(modules.get(module_name), nn.Linear):
            n, k = value.shape
            group = min(64, k)
            codes = rng.integers(0, 3, size=n * k, dtype=np.uint8)
            pad = np.pad(codes, (0, (-len(codes)) % 5), constant_values=1).reshape(-1, 5)
            tensors[name + ".trits"] = mx.array((pad.astype(np.uint16) * [1, 3, 9, 27, 81]).sum(-1).astype(np.uint8))
            tensors[name + ".scales"] = mx.full((n, k // group), .05, dtype=mx.float16)
            tensors[name + ".shape"] = mx.array([n, k], dtype=mx.int32)
            kinds[name], groups[name] = "tern", group
        elif name == "language_model.model.embed_tokens.weight":
            w, s, b = mx.quantize(value.astype(mx.bfloat16), group_size=64, bits=4)
            tensors[name + ".q"], tensors[name + ".scales"], tensors[name + ".biases"] = w, s, b
        else:
            tensors[name] = value.astype(mx.bfloat16)
    packed = tmp_path / "tiny.safetensors"
    mx.save_safetensors(str(packed), tensors, metadata={"format": "kev-ternary-v1", "rot": "none", "emb_bits": "4", "seed": "0", "run": str(source), "kinds": json.dumps(kinds), "groups": json.dumps(groups)})
    output = tmp_path / "bundle"
    build(packed, source, source / "config.json", output)
    return output


def test_offline_standalone_and_strict_loading(tmp_path, monkeypatch):
    bundle = tiny_bundle(tmp_path)

    def forbidden(*args, **kwargs):
        raise AssertionError("original model loader was called")

    import kev.mlx_model
    monkeypatch.setattr(kev.mlx_model, "load_base", forbidden)
    monkeypatch.setattr(kev.mlx_model, "load_full", forbidden)
    monkeypatch.setattr(kev.mlx_model, "merge_lora", forbidden)
    monkeypatch.setattr("huggingface_hub.snapshot_download", forbidden)
    checkpoint = Checkpoint(str(bundle))
    assert checkpoint.packed and not checkpoint.full and checkpoint.hybrid_base()
    tok, model = checkpoint.load("mps", LoadOptions(backend="auto"))
    assert model.dtype == "bfloat16" and model.head.q.in_features == 64 and model.head.temperature == 2.4
    hidden = model._hidden([[tok.convert_tokens_to_ids("billing"), tok.convert_tokens_to_ids("shipping")]])
    assert hidden.shape == (1, 2, 64) and np.isfinite(np.asarray(hidden.astype(mx.float32))).all()
    from fastapi.testclient import TestClient
    from kev.serve import Server, app
    server = Server(checkpoint, tok, model, "mps")
    monkeypatch.setattr(app.state, "server", server, raising=False)
    try:
        with TestClient(app) as client:
            response = client.post("/v1/systemone", json={"state": "charged twice", "questions": {
                "team": {"type": "choice", "instructions": "billing shipping", "criteria": {"billing": "billing", "shipping": "shipping"}}}})
            assert response.status_code == 200
            assert "answers" in response.json() and "usage" in response.json()
            assert client.get("/v1/models").json()["models"][0]["backend"] == "mlx"
    finally:
        server.close()
    for name, module in model.lm.named_modules():
        if isinstance(module, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
            assert module.weight.dtype == mx.uint32
    assert len(checkpoint.weights_sha256()) == 64
    with pytest.raises(ValueError, match="MLX backend"):
        checkpoint.load("cpu", LoadOptions(backend="torch"))
    weights = mx.load(str(bundle / "model.safetensors"))
    mx.eval(weights)
    weights.pop("language_model.model.layers.0.input_layernorm.weight")
    mx.save_safetensors(str(bundle / "model.safetensors"), weights)
    with pytest.raises(ValueError, match="Missing"):
        checkpoint.load("mps", LoadOptions(backend="mlx"))