#!/usr/bin/env python3
"""Immutable source manifest and deterministic DSpark GGUF tensor recipe."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator


REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
SOURCE_MANIFEST_PATH = REPOSITORY_ROOT / "manifest" / "source.json"
with SOURCE_MANIFEST_PATH.open(encoding="utf-8") as _manifest_handle:
    _SOURCE_MANIFEST = json.load(_manifest_handle)

SOURCE_REPOSITORY = str(_SOURCE_MANIFEST["repository"])
SOURCE_REVISION = str(_SOURCE_MANIFEST["revision"])
SOURCE_FILES = {
    str(name): {"size": int(properties["size"]), "sha256": str(properties["sha256"])}
    for name, properties in _SOURCE_MANIFEST["files"].items()
}
if sum(int(properties["size"]) for properties in SOURCE_FILES.values()) != int(
    _SOURCE_MANIFEST["total_size"]
):
    raise ValueError(
        f"{SOURCE_MANIFEST_PATH}: total_size does not match its file entries"
    )

ARCHITECTURE = "deepseek_v4_flash_dspark_draft"
GENERAL_NAME = "DeepSeek-V4-Flash-0731-DSpark-Drafter"
RECIPE_VERSION = 1

BLOCK_SIZE = 5
MARKOV_RANK = 256
NOISE_TOKEN_ID = 128_799
TARGET_LAYER_IDS = (40, 41, 42)
LAYER_COUNT = 3
HEAD_LAYER = 2
N_EMBD = 4096
N_EXPERTS = 256
VOCAB_SIZE = 129_280

GGML_F32 = 0
GGML_F16 = 1
GGML_Q8_0 = 8
GGML_Q2_K = 10
GGML_MXFP4 = 39
GGML_TYPE_NAMES = {
    GGML_F32: "F32",
    GGML_F16: "F16",
    GGML_Q8_0: "Q8_0",
    GGML_Q2_K: "Q2_K",
    GGML_MXFP4: "MXFP4",
}

KIND_PLAIN_F32 = "plain_f32"
KIND_RELABEL_F32 = "relabel_f32"
KIND_FP8_TO_Q8_0 = "fp8_to_q8_0"
KIND_FUSED_EXPERTS_MXFP4 = "fused_experts_mxfp4"
KIND_FUSED_EXPERTS_Q2_K = "fused_experts_q2_k"
KIND_MARKOV_F16 = "markov_f16"

RECIPE_MXFP4_Q8_0 = "mxfp4-q8_0"
RECIPE_Q2_K_Q8_0 = "q2_k-q8_0"


@dataclass(frozen=True)
class ArtifactRecipe:
    """One output policy over the shared canonical tensor/source plan."""

    name: str
    output_filename: str
    routed_expert_kind: str
    routed_expert_type: int
    expected_tensor_bytes: int
    expected_file_size: int
    expected_type_counts: tuple[tuple[int, int], ...]
    q2_k_error_limit: float | None = None

    def manifest_filename(self, kind: str) -> str:
        if kind not in {"build", "validation"}:
            raise ValueError(f"unknown manifest kind: {kind}")
        if self.name == RECIPE_MXFP4_Q8_0:
            return f"{kind}.json"
        return f"{kind}-{self.name}.json"


MXFP4_Q8_0_RECIPE = ArtifactRecipe(
    name=RECIPE_MXFP4_Q8_0,
    output_filename="DeepSeek-V4-Flash-0731-DSpark-Drafter-MXFP4-Q8_0.gguf",
    routed_expert_kind=KIND_FUSED_EXPERTS_MXFP4,
    routed_expert_type=GGML_MXFP4,
    expected_tensor_bytes=10_897_104_284,
    expected_file_size=10_897_110_272,
    expected_type_counts=(
        (GGML_F32, 45),
        (GGML_F16, 2),
        (GGML_Q8_0, 25),
        (GGML_MXFP4, 9),
    ),
)
Q2_K_Q8_0_RECIPE = ArtifactRecipe(
    name=RECIPE_Q2_K_Q8_0,
    output_filename="DeepSeek-V4-Flash-0731-DSpark-Drafter-Q2_K-Q8_0.gguf",
    routed_expert_kind=KIND_FUSED_EXPERTS_Q2_K,
    routed_expert_type=GGML_Q2_K,
    expected_tensor_bytes=6_971_235_740,
    expected_file_size=6_971_241_728,
    expected_type_counts=(
        (GGML_F32, 45),
        (GGML_F16, 2),
        (GGML_Q8_0, 25),
        (GGML_Q2_K, 9),
    ),
    q2_k_error_limit=0.50,
)
RECIPES = {recipe.name: recipe for recipe in (MXFP4_Q8_0_RECIPE, Q2_K_Q8_0_RECIPE)}
DEFAULT_RECIPE = MXFP4_Q8_0_RECIPE
OUTPUT_FILENAME = DEFAULT_RECIPE.output_filename


def resolve_recipe(recipe: ArtifactRecipe | str | None = None) -> ArtifactRecipe:
    if recipe is None:
        return DEFAULT_RECIPE
    if isinstance(recipe, ArtifactRecipe):
        return recipe
    try:
        return RECIPES[recipe]
    except KeyError as error:
        choices = ", ".join(sorted(RECIPES))
        raise ValueError(
            f"unknown recipe {recipe!r}; choose one of: {choices}"
        ) from error


@dataclass(frozen=True)
class TensorRecipe:
    out_name: str
    kind: str
    dims: tuple[int, ...]
    ggml_type: int
    src_name: str = ""
    src_is_bf16: bool = False
    src_shape: tuple[int, ...] = ()
    weight_name: str = ""
    scale_name: str = ""
    rows: int = 0
    cols: int = 0
    name_prefix: str = ""
    expert_kind: str = ""

    @property
    def elements(self) -> int:
        value = 1
        for dim in self.dims:
            value *= dim
        return value

    @property
    def byte_len(self) -> int:
        if self.ggml_type == GGML_F32:
            return self.elements * 4
        if self.ggml_type == GGML_F16:
            return self.elements * 2
        if self.ggml_type == GGML_Q8_0:
            if self.elements % 32:
                raise ValueError(
                    f"{self.out_name}: Q8_0 element count is not block aligned"
                )
            return self.elements // 32 * 34
        if self.ggml_type == GGML_Q2_K:
            if self.elements % 256:
                raise ValueError(
                    f"{self.out_name}: Q2_K element count is not block aligned"
                )
            return self.elements // 256 * 84
        if self.ggml_type == GGML_MXFP4:
            if self.elements % 32:
                raise ValueError(
                    f"{self.out_name}: MXFP4 element count is not block aligned"
                )
            return self.elements // 32 * 17
        raise ValueError(f"{self.out_name}: unsupported GGML type {self.ggml_type}")


def _plain(
    out_name: str,
    src_name: str,
    *,
    bf16: bool,
    dims: tuple[int, ...],
    src_shape: tuple[int, ...] | None = None,
) -> TensorRecipe:
    return TensorRecipe(
        out_name=out_name,
        kind=KIND_PLAIN_F32,
        dims=dims,
        ggml_type=GGML_F32,
        src_name=src_name,
        src_is_bf16=bf16,
        src_shape=dims if src_shape is None else src_shape,
    )


def _fp8(
    out_name: str,
    weight_name: str,
    scale_name: str,
    rows: int,
    cols: int,
) -> TensorRecipe:
    return TensorRecipe(
        out_name=out_name,
        kind=KIND_FP8_TO_Q8_0,
        dims=(cols, rows),
        ggml_type=GGML_Q8_0,
        weight_name=weight_name,
        scale_name=scale_name,
        rows=rows,
        cols=cols,
    )


def _layer_plan(layer: int, artifact_recipe: ArtifactRecipe) -> list[TensorRecipe]:
    prefix = f"mtp.{layer}"
    out: list[TensorRecipe] = [
        _plain(
            f"dspark.{layer}.attn_sinks.weight",
            f"{prefix}.attn.attn_sink",
            bf16=False,
            dims=(64,),
        ),
        _plain(
            f"dspark.{layer}.attn_norm.weight",
            f"{prefix}.attn_norm.weight",
            bf16=True,
            dims=(N_EMBD,),
        ),
        _plain(
            f"dspark.{layer}.ffn_norm.weight",
            f"{prefix}.ffn_norm.weight",
            bf16=True,
            dims=(N_EMBD,),
        ),
        _plain(
            f"dspark.{layer}.attn_kv_a_norm.weight",
            f"{prefix}.attn.kv_norm.weight",
            bf16=True,
            dims=(512,),
        ),
        _plain(
            f"dspark.{layer}.attn_q_a_norm.weight",
            f"{prefix}.attn.q_norm.weight",
            bf16=True,
            dims=(1024,),
        ),
        _fp8(
            f"dspark.{layer}.attn_kv.weight",
            f"{prefix}.attn.wkv.weight",
            f"{prefix}.attn.wkv.scale",
            512,
            N_EMBD,
        ),
        _fp8(
            f"dspark.{layer}.attn_q_a.weight",
            f"{prefix}.attn.wq_a.weight",
            f"{prefix}.attn.wq_a.scale",
            1024,
            N_EMBD,
        ),
        _fp8(
            f"dspark.{layer}.attn_q_b.weight",
            f"{prefix}.attn.wq_b.weight",
            f"{prefix}.attn.wq_b.scale",
            32768,
            1024,
        ),
        _fp8(
            f"dspark.{layer}.attn_output_a.weight",
            f"{prefix}.attn.wo_a.weight",
            f"{prefix}.attn.wo_a.scale",
            8192,
            N_EMBD,
        ),
        _fp8(
            f"dspark.{layer}.attn_output_b.weight",
            f"{prefix}.attn.wo_b.weight",
            f"{prefix}.attn.wo_b.scale",
            N_EMBD,
            8192,
        ),
        _fp8(
            f"dspark.{layer}.ffn_gate_shexp.weight",
            f"{prefix}.ffn.shared_experts.w1.weight",
            f"{prefix}.ffn.shared_experts.w1.scale",
            2048,
            N_EMBD,
        ),
        _fp8(
            f"dspark.{layer}.ffn_up_shexp.weight",
            f"{prefix}.ffn.shared_experts.w3.weight",
            f"{prefix}.ffn.shared_experts.w3.scale",
            2048,
            N_EMBD,
        ),
        _fp8(
            f"dspark.{layer}.ffn_down_shexp.weight",
            f"{prefix}.ffn.shared_experts.w2.weight",
            f"{prefix}.ffn.shared_experts.w2.scale",
            N_EMBD,
            2048,
        ),
        _plain(
            f"dspark.{layer}.ffn_gate_inp.weight",
            f"{prefix}.ffn.gate.weight",
            bf16=True,
            dims=(N_EMBD, N_EXPERTS),
            src_shape=(N_EXPERTS, N_EMBD),
        ),
        _plain(
            f"dspark.{layer}.exp_probs_b.bias",
            f"{prefix}.ffn.gate.bias",
            bf16=False,
            dims=(N_EXPERTS,),
        ),
    ]

    for expert_kind, output_suffix, rows, cols in (
        ("w1", "ffn_gate_exps", 2048, N_EMBD),
        ("w3", "ffn_up_exps", 2048, N_EMBD),
        ("w2", "ffn_down_exps", N_EMBD, 2048),
    ):
        out.append(
            TensorRecipe(
                out_name=f"dspark.{layer}.{output_suffix}.weight",
                kind=artifact_recipe.routed_expert_kind,
                dims=(cols, rows, N_EXPERTS),
                ggml_type=artifact_recipe.routed_expert_type,
                rows=rows,
                cols=cols,
                name_prefix=prefix,
                expert_kind=expert_kind,
            )
        )

    for suffix in ("hc_attn_fn", "hc_ffn_fn"):
        out.append(
            TensorRecipe(
                out_name=f"dspark.{layer}.{suffix}.weight",
                kind=KIND_RELABEL_F32,
                dims=(16384, 24),
                ggml_type=GGML_F32,
                src_name=f"{prefix}.{suffix}",
                rows=24,
                cols=16384,
            )
        )
    for suffix in ("hc_attn_base", "hc_ffn_base"):
        out.append(
            _plain(
                f"dspark.{layer}.{suffix}.weight",
                f"{prefix}.{suffix}",
                bf16=False,
                dims=(24,),
            )
        )
    for suffix in ("hc_attn_scale", "hc_ffn_scale"):
        out.append(
            _plain(
                f"dspark.{layer}.{suffix}.weight",
                f"{prefix}.{suffix}",
                bf16=False,
                dims=(3,),
            )
        )
    return out


def _global_plan() -> list[TensorRecipe]:
    head = f"mtp.{HEAD_LAYER}"
    out = [
        _fp8(
            "dspark.main_proj.weight",
            "mtp.0.main_proj.weight",
            "mtp.0.main_proj.scale",
            N_EMBD,
            12288,
        ),
        _plain(
            "dspark.main_norm.weight",
            "mtp.0.main_norm.weight",
            bf16=True,
            dims=(N_EMBD,),
        ),
        _plain(
            "dspark.norm.weight",
            f"{head}.norm.weight",
            bf16=True,
            dims=(N_EMBD,),
        ),
    ]
    for suffix in ("markov_w1", "markov_w2"):
        out.append(
            TensorRecipe(
                out_name=f"dspark.{suffix}.weight",
                kind=KIND_MARKOV_F16,
                dims=(MARKOV_RANK, VOCAB_SIZE),
                ggml_type=GGML_F16,
                src_name=f"{head}.markov_head.{suffix}.weight",
            )
        )
    out.extend(
        [
            TensorRecipe(
                out_name="dspark.hc_head_fn.weight",
                kind=KIND_RELABEL_F32,
                dims=(16384, 4),
                ggml_type=GGML_F32,
                src_name=f"{head}.hc_head_fn",
                rows=4,
                cols=16384,
            ),
            _plain(
                "dspark.hc_head_base.weight",
                f"{head}.hc_head_base",
                bf16=False,
                dims=(4,),
            ),
            _plain(
                "dspark.hc_head_scale.weight",
                f"{head}.hc_head_scale",
                bf16=False,
                dims=(1,),
            ),
            _plain(
                "dspark.confidence_head.weight",
                f"{head}.confidence_head.proj.weight",
                bf16=True,
                dims=(4352,),
                src_shape=(1, 4352),
            ),
        ]
    )
    return out


def build_plan(
    recipe: ArtifactRecipe | str | None = None,
) -> tuple[TensorRecipe, ...]:
    artifact_recipe = resolve_recipe(recipe)
    tensors: list[TensorRecipe] = []
    for layer in range(LAYER_COUNT):
        tensors.extend(_layer_plan(layer, artifact_recipe))
    tensors.extend(_global_plan())
    if len(tensors) != 81:
        raise AssertionError(f"recipe contains {len(tensors)} tensors, expected 81")
    if len({tensor.out_name for tensor in tensors}) != len(tensors):
        raise AssertionError("recipe contains duplicate output tensor names")
    tensor_bytes = sum(tensor.byte_len for tensor in tensors)
    if tensor_bytes != artifact_recipe.expected_tensor_bytes:
        raise AssertionError(
            f"{artifact_recipe.name}: tensor payload is {tensor_bytes}, "
            f"expected {artifact_recipe.expected_tensor_bytes}"
        )
    type_counts = tuple(
        (ggml_type, sum(tensor.ggml_type == ggml_type for tensor in tensors))
        for ggml_type, _expected in artifact_recipe.expected_type_counts
    )
    if type_counts != artifact_recipe.expected_type_counts:
        raise AssertionError(
            f"{artifact_recipe.name}: type inventory {type_counts} does not match "
            f"{artifact_recipe.expected_type_counts}"
        )
    return tuple(tensors)


def source_expectations(
    tensor: TensorRecipe,
) -> Iterator[tuple[str, str, tuple[int, ...]]]:
    if tensor.kind == KIND_PLAIN_F32:
        yield tensor.src_name, "BF16" if tensor.src_is_bf16 else "F32", tensor.src_shape
    elif tensor.kind == KIND_RELABEL_F32:
        yield tensor.src_name, "F32", (tensor.rows, tensor.cols)
    elif tensor.kind == KIND_FP8_TO_Q8_0:
        yield tensor.weight_name, "F8_E4M3", (tensor.rows, tensor.cols)
        yield (
            tensor.scale_name,
            "F8_E8M0",
            ((tensor.rows + 127) // 128, (tensor.cols + 127) // 128),
        )
    elif tensor.kind in {KIND_FUSED_EXPERTS_MXFP4, KIND_FUSED_EXPERTS_Q2_K}:
        for expert in range(N_EXPERTS):
            prefix = f"{tensor.name_prefix}.ffn.experts.{expert}.{tensor.expert_kind}"
            yield f"{prefix}.weight", "I8", (tensor.rows, tensor.cols // 2)
            yield f"{prefix}.scale", "F8_E8M0", (tensor.rows, tensor.cols // 32)
    elif tensor.kind == KIND_MARKOV_F16:
        yield tensor.src_name, "BF16", (VOCAB_SIZE, MARKOV_RANK)
    else:
        raise ValueError(f"unknown recipe kind {tensor.kind}")


def source_paths(source_dir: Path) -> tuple[Path, ...]:
    return tuple(source_dir / name for name in SOURCE_FILES)


GGUF_METADATA = (
    ("general.architecture", "string", ARCHITECTURE),
    ("general.name", "string", GENERAL_NAME),
    ("general.source.url", "string", f"https://huggingface.co/{SOURCE_REPOSITORY}"),
    ("general.source.revision", "string", SOURCE_REVISION),
    ("general.license", "string", "MIT"),
    ("dspark.recipe_version", "u32", RECIPE_VERSION),
    ("dspark.block_size", "u32", BLOCK_SIZE),
    ("dspark.markov_rank", "u32", MARKOV_RANK),
    ("dspark.noise_token_id", "u32", NOISE_TOKEN_ID),
    ("dspark.target_layer_ids", "array_i32", TARGET_LAYER_IDS),
    ("dspark.layer_count", "u32", LAYER_COUNT),
)