from __future__ import annotations

# ruff: noqa: E402

import json
import io
import struct
import sys
import tempfile
import time
import unittest
from pathlib import Path
from unittest import mock

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts"))

import convert as convert_module
from convert import (
    E4M3_TABLE,
    TensorProducer,
    decode_source_mxfp4,
    quantize_q2_k,
    quantize_q8_0,
    repack_mxfp4,
    round_away_from_zero,
)
from formats import GgufReader, SafeTensorSet, write_gguf
from recipe import (
    ARCHITECTURE,
    GGML_Q2_K,
    GGML_F16,
    GGML_F32,
    GGML_MXFP4,
    GGML_Q8_0,
    GGUF_METADATA,
    KIND_PLAIN_F32,
    MXFP4_Q8_0_RECIPE,
    Q2_K_Q8_0_RECIPE,
    TensorRecipe,
    build_plan,
    source_expectations,
)
from reproduce import _update_sha256sums
from verify import decode_q2_k_blocks


class RecipeTests(unittest.TestCase):
    def test_plan_inventory_and_source_coverage(self) -> None:
        mxfp4_plan = build_plan(MXFP4_Q8_0_RECIPE)
        q2_k_plan = build_plan(Q2_K_Q8_0_RECIPE)
        self.assertEqual(len(mxfp4_plan), 81)
        self.assertEqual(len(q2_k_plan), 81)
        mxfp4_counts = {
            kind: sum(tensor.ggml_type == kind for tensor in mxfp4_plan)
            for kind in (GGML_F32, GGML_F16, GGML_Q8_0, GGML_MXFP4)
        }
        self.assertEqual(
            mxfp4_counts,
            {GGML_F32: 45, GGML_F16: 2, GGML_Q8_0: 25, GGML_MXFP4: 9},
        )
        q2_k_counts = {
            kind: sum(tensor.ggml_type == kind for tensor in q2_k_plan)
            for kind in (GGML_F32, GGML_F16, GGML_Q8_0, GGML_Q2_K)
        }
        self.assertEqual(
            q2_k_counts,
            {GGML_F32: 45, GGML_F16: 2, GGML_Q8_0: 25, GGML_Q2_K: 9},
        )
        self.assertEqual(sum(tensor.byte_len for tensor in mxfp4_plan), 10_897_104_284)
        self.assertEqual(sum(tensor.byte_len for tensor in q2_k_plan), 6_971_235_740)
        self.assertEqual(
            tuple((tensor.out_name, tensor.dims) for tensor in q2_k_plan),
            tuple((tensor.out_name, tensor.dims) for tensor in mxfp4_plan),
        )
        inputs = [
            name
            for tensor in mxfp4_plan
            for name, _dtype, _shape in source_expectations(tensor)
        ]
        self.assertEqual(len(inputs), 4_705)
        self.assertEqual(len(set(inputs)), 4_705)
        self.assertEqual(mxfp4_plan[0].out_name, "dspark.0.attn_sinks.weight")
        self.assertEqual(mxfp4_plan[-1].out_name, "dspark.confidence_head.weight")

    def test_0731_metadata_and_recipe_outputs_are_stable(self) -> None:
        self.assertEqual(ARCHITECTURE, "deepseek_v4_flash_dspark_draft")
        self.assertEqual(
            GGUF_METADATA,
            (
                ("general.architecture", "string", "deepseek_v4_flash_dspark_draft"),
                ("general.name", "string", "DeepSeek-V4-Flash-0731-DSpark-Drafter"),
                (
                    "general.source.url",
                    "string",
                    "https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731",
                ),
                (
                    "general.source.revision",
                    "string",
                    "9e165c30e2704aec5d9d593cce3eebd58bbef1cb",
                ),
                ("general.license", "string", "MIT"),
                ("dspark.recipe_version", "u32", 1),
                ("dspark.block_size", "u32", 5),
                ("dspark.markov_rank", "u32", 256),
                ("dspark.noise_token_id", "u32", 128_799),
                ("dspark.target_layer_ids", "array_i32", (40, 41, 42)),
                ("dspark.layer_count", "u32", 3),
            ),
        )
        self.assertEqual(
            MXFP4_Q8_0_RECIPE.output_filename,
            "DeepSeek-V4-Flash-0731-DSpark-Drafter-MXFP4-Q8_0.gguf",
        )
        self.assertEqual(
            Q2_K_Q8_0_RECIPE.output_filename,
            "DeepSeek-V4-Flash-0731-DSpark-Drafter-Q2_K-Q8_0.gguf",
        )
        self.assertEqual(MXFP4_Q8_0_RECIPE.expected_file_size, 10_897_110_272)
        self.assertEqual(Q2_K_Q8_0_RECIPE.expected_file_size, 6_971_241_728)
        self.assertEqual(
            Q2_K_Q8_0_RECIPE.expected_file_size
            - Q2_K_Q8_0_RECIPE.expected_tensor_bytes,
            MXFP4_Q8_0_RECIPE.expected_file_size
            - MXFP4_Q8_0_RECIPE.expected_tensor_bytes,
        )
        self.assertEqual(MXFP4_Q8_0_RECIPE.manifest_filename("build"), "build.json")
        self.assertEqual(
            Q2_K_Q8_0_RECIPE.manifest_filename("build"),
            "build-q2_k-q8_0.json",
        )

    def test_e4m3_known_values(self) -> None:
        self.assertEqual(float(E4M3_TABLE[0x00]), 0.0)
        self.assertEqual(float(E4M3_TABLE[0x01]), 2**-9)
        self.assertEqual(float(E4M3_TABLE[0x38]), 1.0)
        self.assertEqual(float(E4M3_TABLE[0x40]), 2.0)
        self.assertEqual(float(E4M3_TABLE[0xB8]), -1.0)
        self.assertTrue(np.isnan(E4M3_TABLE[0x7F]))

    def test_q8_rounds_half_away_from_zero(self) -> None:
        values = np.zeros(32, dtype=np.float32)
        values[:5] = [127.0, 0.5, 1.5, -0.5, -1.5]
        encoded = quantize_q8_0(values)
        scale = encoded[0, :2].copy().view("<f2")[0]
        codes = encoded[0, 2:].view(np.int8)
        self.assertEqual(float(scale), 1.0)
        np.testing.assert_array_equal(codes[:5], [127, 1, 2, -1, -2])

    def test_q8_rounding_does_not_promote_values_below_half(self) -> None:
        below_half = np.nextafter(np.float32(0.5), np.float32(0.0), dtype=np.float32)
        values = np.array(
            [below_half, -below_half, np.float32(0.5), np.float32(-0.5)],
            dtype=np.float32,
        )
        np.testing.assert_array_equal(
            round_away_from_zero(values),
            np.array([0.0, -0.0, 1.0, -1.0], dtype=np.float32),
        )

    def test_mxfp4_low16_high16_layout_and_nan_rejection(self) -> None:
        codes = np.arange(32, dtype=np.uint8) & np.uint8(0x0F)
        packed = codes[0::2] | (codes[1::2] << np.uint8(4))
        encoded = repack_mxfp4(packed, np.array([127], dtype=np.uint8), 1, 32)
        self.assertEqual(int(encoded[0, 0, 0]), 127)
        expected = codes[:16] | (codes[16:] << np.uint8(4))
        np.testing.assert_array_equal(encoded[0, 0, 1:], expected)
        with self.assertRaisesRegex(ValueError, "0xff"):
            repack_mxfp4(packed, np.array([0xFF], dtype=np.uint8), 1, 32)

    def test_q2_k_matches_ds4_reference_block_and_decodes_independently(self) -> None:
        values = np.empty((1, 256), dtype=np.float32)
        for index in range(256):
            raw = ((index * 37) % 31) - 15
            multiplier = 1 << ((index // 16) % 4)
            values[0, index] = (
                np.float32(0.0)
                if index % 19 == 0
                else np.float32(raw) * np.float32(multiplier) * np.float32(0.125)
            )
        expected = bytes.fromhex(
            "123357ee124467fe124357ef124467ff"
            "9194e9294e8294e4294e8294d4294e42"
            "e4e9ba4e93a4e93a4e9b94e5394e53a4"
            "4e8294e4694d8294e4294e4690d4294d9"
            "3a4e53a4f5394a53a4f5394e5fa7b5fe038193c"
        )
        encoded = quantize_q2_k(values)
        self.assertEqual(encoded.tobytes(), expected)
        self.assertEqual(encoded[0, :16].tobytes(), expected[:16])
        decoded = decode_q2_k_blocks(encoded)
        relative_error = np.max(np.abs(values - decoded)) / np.max(np.abs(values))
        self.assertLess(float(relative_error), Q2_K_Q8_0_RECIPE.q2_k_error_limit)

        batched = quantize_q2_k(np.concatenate((values, -values), axis=0))
        separate = np.concatenate(
            (quantize_q2_k(values), quantize_q2_k(-values)), axis=0
        )
        np.testing.assert_array_equal(batched, separate)

    def test_threaded_q2_experts_are_written_in_canonical_order(self) -> None:
        rows = 1
        cols = 256

        def source_arrays(expert: int) -> tuple[np.ndarray, np.ndarray]:
            codes = (np.arange(cols, dtype=np.uint16) + expert * 3).astype(np.uint8)
            codes &= np.uint8(0x0F)
            packed = codes[0::2] | (codes[1::2] << np.uint8(4))
            scales = np.full(cols // 32, 127 + expert, dtype=np.uint8)
            return packed, scales

        class DelayedSource:
            def array(self, name: str, _dtype: object) -> np.ndarray:
                expert = int(name.split(".experts.", 1)[1].split(".", 1)[0])
                packed, scales = source_arrays(expert)
                if name.endswith(".weight"):
                    # Later experts finish first, exercising ordered futures.
                    time.sleep(0.01 * (3 - expert))
                    return packed
                return scales

        tensor = TensorRecipe(
            out_name="dspark.test.ffn_gate_exps.weight",
            kind=convert_module.KIND_FUSED_EXPERTS_Q2_K,
            dims=(cols, rows, 4),
            ggml_type=GGML_Q2_K,
            rows=rows,
            cols=cols,
            name_prefix="mtp.blocks.40",
            expert_kind="w1",
        )
        output = io.BytesIO()
        with (
            mock.patch.object(convert_module, "N_EXPERTS", 4),
            mock.patch.object(convert_module, "Q2_K_EXPERT_WORKERS", 4),
        ):
            written = TensorProducer(DelayedSource())._fused_experts_q2_k(
                tensor, output
            )

        expected = bytearray()
        for expert in range(4):
            packed, scales = source_arrays(expert)
            decoded = decode_source_mxfp4(packed, scales, rows, cols)
            expected.extend(quantize_q2_k(decoded).tobytes(order="C"))
        self.assertEqual(written, 4 * 84)
        self.assertEqual(output.getvalue(), bytes(expected))

    def test_sha256sums_update_preserves_other_recipe(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "SHA256SUMS"
            legacy_digest = "1" * 64
            q2_digest = "2" * 64
            _update_sha256sums(path, legacy_digest, "legacy.gguf")
            _update_sha256sums(path, q2_digest, "q2.gguf")
            self.assertEqual(
                path.read_text(encoding="ascii"),
                f"{legacy_digest}  legacy.gguf\n{q2_digest}  q2.gguf\n",
            )

    def test_tiny_gguf_round_trip(self) -> None:
        tensor = TensorRecipe(
            out_name="dspark.test.weight",
            kind=KIND_PLAIN_F32,
            dims=(4,),
            ggml_type=GGML_F32,
            src_name="mtp.test",
            src_shape=(4,),
        )
        data = struct.pack("<4f", 1.0, 2.0, 3.0, 4.0)
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "tiny.gguf"

            def produce(_tensor: TensorRecipe, handle: object) -> int:
                return handle.write(data)

            write_gguf(path, (tensor,), produce)
            with GgufReader(path) as reader:
                self.assertEqual(reader.tensors[0].name, tensor.out_name)
                self.assertEqual(reader.tensors[0].dims, tensor.dims)
                self.assertEqual(reader.tensors[0].ggml_type, GGML_F32)
                observed = reader.tensor_array(reader.tensors[0], "<f4").copy()
            np.testing.assert_array_equal(observed, [1.0, 2.0, 3.0, 4.0])
            self.assertEqual(path.stat().st_size % 32, 0)

    def test_safetensors_rejects_bad_offsets(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "bad.safetensors"
            header = {
                "mtp.bad": {
                    "dtype": "F32",
                    "shape": [1],
                    "data_offsets": [0, 8],
                }
            }
            encoded = json.dumps(header, separators=(",", ":")).encode()
            path.write_bytes(struct.pack("<Q", len(encoded)) + encoded + b"\0" * 4)
            with self.assertRaisesRegex(ValueError, "overruns|needs"):
                SafeTensorSet([path])


if __name__ == "__main__":
    unittest.main()