"""K-quant wire format and dequantization tests.

The node advertises `quant_types: [F32, F16, Q8_0, Q4_K, Q6_K]` and the client
examples send Q4_K, so these paths are part of the protocol contract with
llama.cpp: block sizes decide where a matrix ends on the socket, and the bit
layout decides whether the numbers mean anything.

The reference implementations below are transcribed straight from ggml
(`dequantize_row_q4_K`, `dequantize_row_q6_K`, `get_scale_min_k4` in
llama.cpp/ggml-quants.c) and deliberately written as slow scalar loops, so they
can be read against the C side line by line. If the `gguf` package happens to be
installed, one more test cross-checks against its independent decoder.
"""

import struct

import numpy as np
import pytest

import coffer_router
import mlx_matmul_server as server

QK_K = 256
QK8_0 = 32

Q8_0, Q4_K, Q6_K = server.GGML_TYPE_Q8_0, server.GGML_TYPE_Q4_K, server.GGML_TYPE_Q6_K

# ggml-common.h block sizes. Q4_K is 144: `scales[12]` carries the 8 scales *and*
# the 8 mins at six bits each — there is no extra 4-byte mins array.
GGML_BLOCK_BYTES = {Q8_0: 34, Q4_K: 144, Q6_K: 210}
GGML_BLOCK_ELEMENTS = {Q8_0: QK8_0, Q4_K: QK_K, Q6_K: QK_K}


# ─── Block builders (valid random blocks, no quantizer needed) ────────────────

def _rng(seed):
    return np.random.default_rng(seed)


def make_q8_0(n_blocks, seed=1):
    r = _rng(seed)
    out = b""
    for _ in range(n_blocks):
        out += np.float16(r.uniform(0.001, 0.2)).tobytes()
        out += r.integers(-128, 128, QK8_0, dtype=np.int8).tobytes()
    return out


def make_q4_k(n_blocks, seed=2):
    r = _rng(seed)
    out = b""
    for _ in range(n_blocks):
        out += np.float16(r.uniform(0.001, 0.05)).tobytes()      # d
        out += np.float16(r.uniform(0.001, 0.05)).tobytes()      # dmin
        out += r.integers(0, 256, 12, dtype=np.uint8).tobytes()  # packed scales+mins
        out += r.integers(0, 256, 128, dtype=np.uint8).tobytes() # 4-bit quants
    return out


def make_q6_k(n_blocks, seed=3):
    r = _rng(seed)
    out = b""
    for _ in range(n_blocks):
        out += r.integers(0, 256, 128, dtype=np.uint8).tobytes()   # ql
        out += r.integers(0, 256, 64, dtype=np.uint8).tobytes()    # qh
        out += r.integers(-128, 128, 16, dtype=np.int8).tobytes()  # scales
        out += np.float16(r.uniform(0.001, 0.05)).tobytes()        # d
    return out


BUILDERS = {Q8_0: make_q8_0, Q4_K: make_q4_k, Q6_K: make_q6_k}
DEQUANTIZERS = {Q8_0: server.dequant_q8_0, Q4_K: server.dequant_q4_k,
                Q6_K: server.dequant_q6_k}


# ─── Scalar reference, transcribed from ggml ──────────────────────────────────

def ref_q8_0(blob, n_elements):
    out = []
    for b in range(len(blob) // 34):
        blk = blob[b * 34:(b + 1) * 34]
        d = struct.unpack("<e", blk[0:2])[0]
        out += [d * q for q in struct.unpack("<32b", blk[2:34])]
    return np.asarray(out[:n_elements], dtype=np.float32)


def ref_q4_k(blob, n_elements):
    out = []
    for b in range(len(blob) // 144):
        blk = blob[b * 144:(b + 1) * 144]
        d = struct.unpack("<e", blk[0:2])[0]
        dmin = struct.unpack("<e", blk[2:4])[0]
        sc, qs = blk[4:16], blk[16:144]

        def get_scale_min_k4(j):           # ggml-quants.c
            if j < 4:
                return sc[j] & 63, sc[j + 4] & 63
            return ((sc[j + 4] & 0x0F) | ((sc[j - 4] >> 6) << 4),
                    (sc[j + 4] >> 4) | ((sc[j] >> 6) << 4))

        y, is_ = [], 0
        for group in range(4):             # 4 groups of 64 elements
            s1, m1 = get_scale_min_k4(is_)
            s2, m2 = get_scale_min_k4(is_ + 1)
            q = qs[group * 32:(group + 1) * 32]
            y += [d * s1 * (v & 0x0F) - dmin * m1 for v in q]
            y += [d * s2 * (v >> 4) - dmin * m2 for v in q]
            is_ += 2
        out += y
    return np.asarray(out[:n_elements], dtype=np.float32)


def ref_q6_k(blob, n_elements):
    out = []
    for b in range(len(blob) // 210):
        blk = blob[b * 210:(b + 1) * 210]
        ql, qh = blk[0:128], blk[128:192]
        sc = struct.unpack("<16b", blk[192:208])
        d = struct.unpack("<e", blk[208:210])[0]

        y = [0.0] * QK_K
        for half in range(2):              # ggml walks 128 elements at a time
            qlh, qhh, sch = ql[half * 64:], qh[half * 32:], sc[half * 8:]
            base = half * 128
            for l in range(32):
                i_s = l // 16
                q1 = ((qlh[l] & 0x0F) | (((qhh[l] >> 0) & 3) << 4)) - 32
                q2 = ((qlh[l + 32] & 0x0F) | (((qhh[l] >> 2) & 3) << 4)) - 32
                q3 = ((qlh[l] >> 4) | (((qhh[l] >> 4) & 3) << 4)) - 32
                q4 = ((qlh[l + 32] >> 4) | (((qhh[l] >> 6) & 3) << 4)) - 32
                y[base + l] = d * sch[i_s + 0] * q1
                y[base + l + 32] = d * sch[i_s + 2] * q2
                y[base + l + 64] = d * sch[i_s + 4] * q3
                y[base + l + 96] = d * sch[i_s + 6] * q4
        out += y
    return np.asarray(out[:n_elements], dtype=np.float32)


REFERENCES = {Q8_0: ref_q8_0, Q4_K: ref_q4_k, Q6_K: ref_q6_k}


def assert_close_in_fp16(got, ref):
    """Production code returns fp16, the reference fp32 — allow one fp16 ULP."""
    got = np.asarray(got, dtype=np.float32)
    assert got.shape == ref.shape
    tol = max(1e-4, float(np.abs(ref).max()) * 2 ** -10)
    assert np.abs(got - ref).max() <= tol


# ─── Wire sizes ──────────────────────────────────────────────────────────────

@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_block_byte_size_is_the_ggml_one(tensor_type):
    assert server.BLOCK_BYTES[tensor_type] == GGML_BLOCK_BYTES[tensor_type]
    assert server.BLOCK_ELEMENTS[tensor_type] == GGML_BLOCK_ELEMENTS[tensor_type]


def test_q4_k_block_is_144_not_148():
    # Regression sentinel: 148 (an invented 4-byte mins field) made every real
    # Q4_K request read 4 bytes/block too many and desync the connection.
    assert server.get_type_size(Q4_K, QK_K) == 144
    assert coffer_router.get_type_size(Q4_K, QK_K) == 144


@pytest.mark.parametrize("tensor_type", [server.GGML_TYPE_F32, server.GGML_TYPE_F16,
                                         Q8_0, Q4_K, Q6_K])
@pytest.mark.parametrize("elements", [32, 256, 512, 4096, 4096 * 4096])
def test_router_and_server_agree_on_payload_size(tensor_type, elements):
    assert (coffer_router.get_type_size(tensor_type, elements)
            == server.get_type_size(tensor_type, elements))


@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_declared_size_matches_what_a_real_block_stream_occupies(tensor_type):
    n_blocks = 5
    blob = BUILDERS[tensor_type](n_blocks)
    elements = n_blocks * GGML_BLOCK_ELEMENTS[tensor_type]
    assert len(blob) == server.get_type_size(tensor_type, elements)


# ─── Numerics ────────────────────────────────────────────────────────────────

@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
@pytest.mark.parametrize("n_blocks", [1, 2, 9])
def test_dequant_matches_ggml_reference(tensor_type, n_blocks):
    blob = BUILDERS[tensor_type](n_blocks)
    elements = n_blocks * GGML_BLOCK_ELEMENTS[tensor_type]
    got = DEQUANTIZERS[tensor_type](blob, elements)
    assert got.dtype == np.float16
    assert_close_in_fp16(got, REFERENCES[tensor_type](blob, elements))


@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_dequant_of_partial_trailing_block(tensor_type):
    n_blocks = 3
    blob = BUILDERS[tensor_type](n_blocks)
    elements = n_blocks * GGML_BLOCK_ELEMENTS[tensor_type] - 7
    got = DEQUANTIZERS[tensor_type](blob, elements)
    assert len(got) == elements
    assert_close_in_fp16(got, REFERENCES[tensor_type](blob, elements))


@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_dequant_rejects_truncated_payload(tensor_type):
    blob = BUILDERS[tensor_type](2)
    elements = 2 * GGML_BLOCK_ELEMENTS[tensor_type]
    with pytest.raises(ValueError):
        DEQUANTIZERS[tensor_type](blob[:-1], elements)


def test_q6_k_negative_quants_do_not_wrap_around():
    """`q - 32` on a uint8 used to wrap: a quant of 0 came out as +224."""
    ql = bytes([0x00] * 128)           # every low and high nibble = 0
    qh = bytes([0x00] * 64)            # every high bit pair = 0  -> q = 0
    scales = bytes([1] * 16)
    d = np.float16(1.0).tobytes()
    got = np.asarray(server.dequant_q6_k(ql + qh + scales + d, QK_K), dtype=np.float32)
    assert np.all(got == -32.0)


def test_q4_k_sub_block_scales_and_mins_are_applied_per_32_values():
    """Unit scale, zero min: output must be the raw 4-bit values themselves."""
    scales = bytes([1, 1, 1, 1, 0, 0, 0, 0, 0x11, 0x11, 0x11, 0x11])
    qs = bytes(range(256 // 2))
    blob = np.float16(1.0).tobytes() + np.float16(0.0).tobytes() + scales + qs
    got = np.asarray(server.dequant_q4_k(blob, QK_K), dtype=np.float32)
    expected = ref_q4_k(blob, QK_K)
    assert_close_in_fp16(got, expected)
    # first 32 values are the low nibbles of qs[0:32] at scale 1
    assert np.array_equal(got[:32], np.array([b & 0x0F for b in qs[:32]], dtype=np.float32))


@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_dequantize_dispatches_to_the_right_decoder(tensor_type):
    blob = BUILDERS[tensor_type](2)
    elements = 2 * GGML_BLOCK_ELEMENTS[tensor_type]
    assert_close_in_fp16(server.dequantize(blob, tensor_type, elements),
                         REFERENCES[tensor_type](blob, elements))


def test_gguf_cross_check():
    """Optional: verify against the gguf package's independent decoder."""
    gguf_quants = pytest.importorskip("gguf.quants")
    from gguf.constants import GGMLQuantizationType

    mapping = {Q8_0: GGMLQuantizationType.Q8_0,
               Q4_K: GGMLQuantizationType.Q4_K,
               Q6_K: GGMLQuantizationType.Q6_K}
    for tensor_type, gguf_type in mapping.items():
        blob = BUILDERS[tensor_type](4)
        elements = 4 * GGML_BLOCK_ELEMENTS[tensor_type]
        ref = gguf_quants.dequantize(
            np.frombuffer(blob, dtype=np.uint8), gguf_type
        ).astype(np.float32).ravel()
        assert_close_in_fp16(DEQUANTIZERS[tensor_type](blob, elements), ref)


# ─── Framing: a K-quant request must not desync the connection ───────────────

class ScriptedSocket:
    """Feeds handle_client a fixed byte script and collects what it writes."""

    def __init__(self, script: bytes):
        self._inbox = script
        self.sent = bytearray()
        self.closed = False

    def recv(self, n):
        chunk, self._inbox = self._inbox[:n], self._inbox[n:]
        return chunk

    def send(self, data):
        self.sent.extend(data)
        return len(data)

    def close(self):
        self.closed = True


def v3_request(M, N, K, A_type, B_type, A_data, B_data):
    return (struct.pack("<IIIIII", server.MAGIC_V3, M, N, K, A_type, B_type)
            + A_data + B_data)


def parse_responses(sent, expected):
    """Split the server's output into (status, M, N) triples."""
    out, offset = [], 0
    for _ in range(expected):
        magic, status, M, N = struct.unpack("<IIII", sent[offset:offset + 16])
        assert magic == server.MAGIC_V3
        offset += 16 + M * N * 2
        out.append((status, M, N))
    assert offset == len(sent), f"trailing bytes: {len(sent) - offset}"
    return out


@pytest.mark.parametrize("tensor_type", [Q8_0, Q4_K, Q6_K])
def test_two_quantized_requests_on_one_connection_stay_in_sync(tensor_type):
    M, N, K = 1, 1, GGML_BLOCK_ELEMENTS[tensor_type]
    A = BUILDERS[tensor_type](1, seed=11)
    B = BUILDERS[tensor_type](1, seed=12)
    request = v3_request(M, N, K, tensor_type, tensor_type, A, B)

    sock = ScriptedSocket(request * 2)
    server.handle_client(sock, ("test", 0))

    assert sock.closed
    assert parse_responses(bytes(sock.sent), 2) == [(0, M, N), (0, M, N)]


@pytest.mark.parametrize("tensor_type", [Q4_K, Q6_K])
def test_quantized_matmul_over_the_wire_matches_local_numpy(tensor_type):
    M, N, K = 1, 2, QK_K
    A = BUILDERS[tensor_type](1, seed=21)
    B = BUILDERS[tensor_type](2, seed=22)          # K*N = 512 elements = 2 blocks

    sock = ScriptedSocket(v3_request(M, N, K, tensor_type, tensor_type, A, B))
    server.handle_client(sock, ("test", 0))

    status, rM, rN = parse_responses(bytes(sock.sent), 1)[0]
    assert (status, rM, rN) == (0, M, N)
    result = np.frombuffer(bytes(sock.sent)[16:], dtype=np.float16).reshape(rM, rN)

    A_ref = REFERENCES[tensor_type](A, M * K).reshape(M, K)
    B_ref = REFERENCES[tensor_type](B, K * N).reshape(K, N)
    expected = A_ref @ B_ref
    np.testing.assert_allclose(result.astype(np.float32), expected,
                               rtol=2e-2, atol=1e-2)