#!/usr/bin/env python3
"""Independently verify GGUF structure and every tensor against pinned sources."""

from __future__ import annotations

import argparse
import json
import math
import os
from pathlib import Path

import numpy as np

from formats import (
    ALIGNMENT,
    GgufReader,
    SafeTensorSet,
    align_up,
    sha256_path,
    verify_source_files,
)
from recipe import (
    DEFAULT_RECIPE,
    GGML_F16,
    GGML_F32,
    GGML_MXFP4,
    GGML_Q2_K,
    GGML_Q8_0,
    GGML_TYPE_NAMES,
    GGUF_METADATA,
    KIND_PLAIN_F32,
    N_EXPERTS,
    RECIPES,
    SOURCE_REPOSITORY,
    SOURCE_REVISION,
    ArtifactRecipe,
    TensorRecipe,
    build_plan,
    resolve_recipe,
    source_paths,
)


def _decode_e4m3_table() -> np.ndarray:
    values = []
    for byte in range(256):
        sign = -1.0 if byte & 0x80 else 1.0
        exponent = (byte >> 3) & 0x0F
        mantissa = byte & 0x07
        if exponent == 0x0F and mantissa == 0x07:
            value = np.nan
        elif exponent == 0:
            value = sign * (mantissa / 8.0) * math.ldexp(1.0, -6)
        else:
            value = sign * (1.0 + mantissa / 8.0) * math.ldexp(1.0, exponent - 7)
        values.append(np.float32(value))
    return np.array(values, dtype=np.float32)


VERIFY_E4M3 = _decode_e4m3_table()
VERIFY_E8M0 = np.array(
    [
        np.float32(np.nan) if code == 0xFF else np.float32(math.ldexp(1.0, code - 127))
        for code in range(256)
    ],
    dtype=np.float32,
)
VERIFY_FP4 = np.array(
    [
        0.0,
        0.5,
        1.0,
        1.5,
        2.0,
        3.0,
        4.0,
        6.0,
        0.0,
        -0.5,
        -1.0,
        -1.5,
        -2.0,
        -3.0,
        -4.0,
        -6.0,
    ],
    dtype=np.float32,
)
# Independently decode and compare at most one routed expert at a time.
VERIFY_EXPERT_ROW_BATCH = 4096


def decode_q2_k_blocks(encoded: np.ndarray) -> np.ndarray:
    """Independently decode GGML Q2_K's 84-byte/256-value block layout."""

    blocks = np.asarray(encoded, dtype=np.uint8)
    if blocks.ndim != 2 or blocks.shape[1] != 84:
        raise ValueError("Q2_K payload must have shape (n, 84)")
    scale_bytes = blocks[:, :16]
    d = blocks[:, 80:82].copy().reshape(-1).view("<f2").astype(np.float32)
    dmin = blocks[:, 82:84].copy().reshape(-1).view("<f2").astype(np.float32)
    if not np.isfinite(d).all() or not np.isfinite(dmin).all():
        raise ValueError("Q2_K payload contains a non-finite block scale")

    codes = np.empty((blocks.shape[0], 256), dtype=np.uint8)
    for start in (0, 128):
        offset = 16 + start // 4
        packed = blocks[:, offset : offset + 32]
        codes[:, start : start + 32] = packed & np.uint8(0x03)
        codes[:, start + 32 : start + 64] = (packed >> np.uint8(2)) & np.uint8(0x03)
        codes[:, start + 64 : start + 96] = (packed >> np.uint8(4)) & np.uint8(0x03)
        codes[:, start + 96 : start + 128] = packed >> np.uint8(6)

    group_scales = np.multiply(
        d[:, None], scale_bytes & np.uint8(0x0F), dtype=np.float32
    )
    group_mins = np.multiply(
        dmin[:, None], scale_bytes >> np.uint8(4), dtype=np.float32
    )
    decoded = np.subtract(
        np.multiply(
            group_scales[..., None],
            codes.reshape(-1, 16, 16),
            dtype=np.float32,
        ),
        group_mins[..., None],
        dtype=np.float32,
    )
    return decoded.reshape(-1, 256)


def _decode_source_mxfp4_rows(
    packed: np.ndarray,
    scale_codes: np.ndarray,
    rows: int,
    cols: int,
) -> np.ndarray:
    if cols % 32:
        raise ValueError("source MXFP4 column count must be divisible by 32")
    block_count = cols // 32
    source = np.asarray(packed, dtype=np.uint8).reshape(rows, block_count, 16)
    scales = np.asarray(scale_codes, dtype=np.uint8).reshape(rows, block_count)
    if np.any(scales == 0xFF):
        raise ValueError("source MXFP4 contains reserved E8M0 NaN scale 0xff")
    codes = np.empty((rows, block_count, 32), dtype=np.uint8)
    codes[..., 0::2] = source & np.uint8(0x0F)
    codes[..., 1::2] = source >> np.uint8(4)
    decoded = VERIFY_FP4[codes]
    decoded *= VERIFY_E8M0[scales][..., None]
    return decoded.reshape(rows, cols)


def _expected_metadata() -> dict[str, object]:
    expected: dict[str, object] = {}
    for key, kind, value in GGUF_METADATA:
        if kind == "array_i32":
            expected[key] = tuple(int(item) for item in value)
        elif kind == "u32":
            expected[key] = int(value)
        else:
            expected[key] = str(value)
    return expected


def _assert_equal(name: str, observed: np.ndarray, expected: np.ndarray) -> None:
    if observed.shape != expected.shape:
        raise ValueError(
            f"{name}: expected shape {expected.shape}, found {observed.shape}"
        )
    if not np.array_equal(observed, expected):
        mismatch = np.flatnonzero(observed.reshape(-1) != expected.reshape(-1))
        first = int(mismatch[0]) if mismatch.size else -1
        raise ValueError(f"{name}: exact comparison failed at element {first}")


def _verify_f32(
    source: SafeTensorSet,
    reader: GgufReader,
    tensor: TensorRecipe,
) -> None:
    output = reader.tensor_array(reader.tensor(tensor.out_name), "<u4")
    if tensor.kind == KIND_PLAIN_F32 and tensor.src_is_bf16:
        bits = source.array(tensor.src_name, "<u2").reshape(-1)
        expected = bits.astype("<u4") << np.uint32(16)
    else:
        expected = source.array(tensor.src_name, "<u4").reshape(-1)
    _assert_equal(tensor.out_name, output, expected)


def _verify_markov_f16(
    source: SafeTensorSet,
    reader: GgufReader,
    tensor: TensorRecipe,
) -> None:
    output = reader.tensor_array(reader.tensor(tensor.out_name), "<u2")
    bits = source.array(tensor.src_name, "<u2").reshape(-1)
    chunk = 8 * 1024 * 1024
    for start in range(0, bits.size, chunk):
        stop = min(start + chunk, bits.size)
        f32_bits = bits[start:stop].astype("<u4") << np.uint32(16)
        expected = f32_bits.view("<f4").astype("<f2").view("<u2")
        _assert_equal(
            f"{tensor.out_name}[{start}:{stop}]",
            output[start:stop],
            expected,
        )


def _verify_q8_0(
    source: SafeTensorSet,
    reader: GgufReader,
    tensor: TensorRecipe,
) -> float:
    info = reader.tensor(tensor.out_name)
    encoded = reader.tensor_array(info, np.uint8).reshape(-1, 34)
    weights = source.array(tensor.weight_name, np.uint8).reshape(
        tensor.rows, tensor.cols
    )
    scale_codes = source.array(tensor.scale_name, np.uint8).reshape(
        (tensor.rows + 127) // 128,
        (tensor.cols + 127) // 128,
    )
    max_relative_error = np.float32(0.0)
    block_cursor = 0
    for block_row in range(scale_codes.shape[0]):
        row_start = block_row * 128
        row_end = min(row_start + 128, tensor.rows)
        reference = VERIFY_E4M3[weights[row_start:row_end]]
        column_scales = np.repeat(VERIFY_E8M0[scale_codes[block_row]], 128)[
            : tensor.cols
        ]
        if not np.isfinite(reference).all() or not np.isfinite(column_scales).all():
            raise ValueError(f"{tensor.out_name}: non-finite pinned FP8 source")
        reference *= column_scales[None, :]
        reference_blocks = reference.reshape(-1, 32)
        block_count = reference_blocks.shape[0]
        output_blocks = encoded[block_cursor : block_cursor + block_count]
        maxima = np.max(np.abs(reference_blocks), axis=1).astype(np.float32, copy=False)
        expected_scales = maxima / np.float32(127.0)
        expected_scale_bits = expected_scales.astype("<f2").view("<u2")
        observed_scale_bits = output_blocks[:, :2].copy().reshape(-1).view("<u2")
        _assert_equal(
            f"{tensor.out_name}/scale-bits-{block_cursor}",
            observed_scale_bits,
            expected_scale_bits,
        )
        inverse = np.zeros_like(expected_scales)
        np.divide(
            np.float32(1.0),
            expected_scales,
            out=inverse,
            where=expected_scales != 0,
        )
        normalized = reference_blocks * inverse[:, None]
        truncated = np.trunc(normalized).astype(np.float32, copy=False)
        fraction = np.abs(normalized - truncated)
        expected_codes = np.where(
            fraction >= np.float32(0.5),
            truncated + np.copysign(np.float32(1.0), normalized),
            truncated,
        ).astype(np.int8)
        observed_codes = output_blocks[:, 2:].view(np.int8)
        _assert_equal(
            f"{tensor.out_name}/codes-{block_cursor}",
            observed_codes,
            expected_codes,
        )
        stored_scale = (
            output_blocks[:, :2].copy().reshape(-1).view("<f2").astype(np.float32)
        )
        if not np.isfinite(stored_scale).all():
            raise ValueError(f"{tensor.out_name}: non-finite Q8_0 scale")
        codes = output_blocks[:, 2:].view(np.int8).astype(np.float32)
        decoded = stored_scale[:, None] * codes
        denominator = np.maximum(
            np.max(np.abs(reference_blocks), axis=1), np.float32(1e-6)
        )
        relative = np.abs(reference_blocks - decoded) / denominator[:, None]
        if not np.isfinite(relative).all():
            raise ValueError(f"{tensor.out_name}: non-finite Q8_0 error")
        max_relative_error = np.maximum(max_relative_error, np.max(relative))
        block_cursor += block_count
    if block_cursor != encoded.shape[0]:
        raise ValueError(f"{tensor.out_name}: Q8_0 block count mismatch")
    result = float(max_relative_error)
    if result > 0.01:
        raise ValueError(
            f"{tensor.out_name}: Q8_0 max relative error {result:.6f} > 0.01"
        )
    return result


def _expected_mxfp4_codes(packed: np.ndarray, rows: int, cols: int) -> np.ndarray:
    blocks = cols // 32
    packed_blocks = np.asarray(packed, dtype=np.uint8).reshape(rows, blocks, 16)
    low = np.empty((rows, blocks, 16), dtype=np.uint8)
    high = np.empty((rows, blocks, 16), dtype=np.uint8)
    low[..., 0::2] = packed_blocks[..., :8] & np.uint8(0x0F)
    low[..., 1::2] = packed_blocks[..., :8] >> np.uint8(4)
    high[..., 0::2] = packed_blocks[..., 8:] & np.uint8(0x0F)
    high[..., 1::2] = packed_blocks[..., 8:] >> np.uint8(4)
    return low | (high << np.uint8(4))


def _verify_mxfp4(
    source: SafeTensorSet,
    reader: GgufReader,
    tensor: TensorRecipe,
) -> None:
    blocks = tensor.cols // 32
    output = reader.tensor_array(reader.tensor(tensor.out_name), np.uint8).reshape(
        N_EXPERTS, tensor.rows, blocks, 17
    )
    for expert in range(N_EXPERTS):
        prefix = f"{tensor.name_prefix}.ffn.experts.{expert}.{tensor.expert_kind}"
        packed = source.array(f"{prefix}.weight", np.uint8)
        scales = source.array(f"{prefix}.scale", np.uint8).reshape(tensor.rows, blocks)
        if np.any(scales == 0xFF):
            raise ValueError(f"{prefix}.scale: contains reserved E8M0 NaN")
        _assert_equal(
            f"{tensor.out_name}/expert-{expert}/scale",
            output[expert, ..., 0],
            scales,
        )
        expected_codes = _expected_mxfp4_codes(packed, tensor.rows, tensor.cols)
        _assert_equal(
            f"{tensor.out_name}/expert-{expert}/codes",
            output[expert, ..., 1:],
            expected_codes,
        )


def _verify_q2_k(
    source: SafeTensorSet,
    reader: GgufReader,
    tensor: TensorRecipe,
    error_limit: float,
) -> float:
    blocks_per_row = tensor.cols // 256
    output = reader.tensor_array(reader.tensor(tensor.out_name), np.uint8).reshape(
        N_EXPERTS, tensor.rows, blocks_per_row, 84
    )
    max_relative_error = np.float32(0.0)
    source_scale_blocks = tensor.cols // 32
    for expert in range(N_EXPERTS):
        prefix = f"{tensor.name_prefix}.ffn.experts.{expert}.{tensor.expert_kind}"
        packed = source.array(f"{prefix}.weight", np.uint8).reshape(
            tensor.rows, tensor.cols // 2
        )
        scales = source.array(f"{prefix}.scale", np.uint8).reshape(
            tensor.rows, source_scale_blocks
        )
        for row_start in range(0, tensor.rows, VERIFY_EXPERT_ROW_BATCH):
            row_end = min(row_start + VERIFY_EXPERT_ROW_BATCH, tensor.rows)
            row_count = row_end - row_start
            reference = _decode_source_mxfp4_rows(
                packed[row_start:row_end],
                scales[row_start:row_end],
                row_count,
                tensor.cols,
            ).reshape(-1, 256)
            encoded = output[expert, row_start:row_end].reshape(-1, 84)
            decoded = decode_q2_k_blocks(encoded)
            if not np.isfinite(reference).all() or not np.isfinite(decoded).all():
                raise ValueError(f"{tensor.out_name}: non-finite Q2_K comparison")
            denominator = np.max(np.abs(reference), axis=1)
            absolute_error = np.max(np.abs(reference - decoded), axis=1)
            relative = np.zeros_like(absolute_error)
            np.divide(
                absolute_error,
                denominator,
                out=relative,
                where=denominator != 0,
            )
            relative[np.logical_and(denominator == 0, absolute_error != 0)] = np.inf
            max_relative_error = np.maximum(max_relative_error, np.max(relative))
    result = float(max_relative_error)
    if result > error_limit:
        raise ValueError(
            f"{tensor.out_name}: Q2_K max block-relative error "
            f"{result:.6f} > {error_limit:.6f}"
        )
    return result


def _verify_zero_padding(reader: GgufReader) -> None:
    header_padding = reader.bytes_at(reader.directory_end, reader.data_start)
    if any(header_padding):
        raise ValueError("non-zero GGUF padding before tensor data")
    tensors = reader.tensors
    for index, tensor in enumerate(tensors):
        data_end = tensor.absolute_offset + tensor.byte_len
        next_start = (
            tensors[index + 1].absolute_offset
            if index + 1 < len(tensors)
            else reader.data_start + align_up(tensor.relative_offset + tensor.byte_len)
        )
        padding = reader.bytes_at(data_end, next_start)
        if any(padding):
            raise ValueError(f"non-zero GGUF alignment padding after {tensor.name}")


def verify(
    source_dir: Path,
    gguf_path: Path,
    *,
    recipe: ArtifactRecipe | str | None = None,
) -> dict[str, object]:
    artifact_recipe = resolve_recipe(recipe)
    source_dir = source_dir.resolve()
    gguf_path = gguf_path.resolve()
    print("[verify] validating pinned source sizes and SHA-256 values", flush=True)
    sources = verify_source_files(source_dir)
    plan = build_plan(artifact_recipe)
    max_q8_error = 0.0
    max_q2_k_error = 0.0
    exact_f32 = 0
    exact_f16 = 0
    exact_mxfp4 = 0
    q2_k_tensors = 0
    q8_tensors = 0
    with (
        SafeTensorSet(source_paths(source_dir)) as source,
        GgufReader(gguf_path) as reader,
    ):
        source.validate_plan(plan)
        expected_metadata = _expected_metadata()
        if reader.metadata != expected_metadata:
            raise ValueError(
                f"GGUF metadata mismatch: expected {expected_metadata}, found {reader.metadata}"
            )
        observed_names = tuple(tensor.name for tensor in reader.tensors)
        expected_names = tuple(tensor.out_name for tensor in plan)
        if observed_names != expected_names:
            raise ValueError("GGUF tensor order or names do not match the recipe")
        expected_relative_offset = 0
        for recipe_tensor, gguf_tensor in zip(plan, reader.tensors, strict=True):
            if gguf_tensor.relative_offset != expected_relative_offset:
                raise ValueError(
                    f"{recipe_tensor.out_name}: expected relative offset "
                    f"{expected_relative_offset}, found {gguf_tensor.relative_offset}"
                )
            expected_relative_offset = align_up(
                expected_relative_offset + recipe_tensor.byte_len
            )
            if gguf_tensor.dims != recipe_tensor.dims:
                raise ValueError(
                    f"{recipe_tensor.out_name}: expected dims {recipe_tensor.dims}, "
                    f"found {gguf_tensor.dims}"
                )
            if gguf_tensor.ggml_type != recipe_tensor.ggml_type:
                raise ValueError(
                    f"{recipe_tensor.out_name}: expected type {recipe_tensor.ggml_type}, "
                    f"found {gguf_tensor.ggml_type}"
                )
            print(f"[verify] {recipe_tensor.out_name}", flush=True)
            if recipe_tensor.ggml_type == GGML_F32:
                _verify_f32(source, reader, recipe_tensor)
                exact_f32 += 1
            elif recipe_tensor.ggml_type == GGML_F16:
                _verify_markov_f16(source, reader, recipe_tensor)
                exact_f16 += 1
            elif recipe_tensor.ggml_type == GGML_Q8_0:
                q8_tensors += 1
                max_q8_error = max(
                    max_q8_error, _verify_q8_0(source, reader, recipe_tensor)
                )
            elif recipe_tensor.ggml_type == GGML_MXFP4:
                _verify_mxfp4(source, reader, recipe_tensor)
                exact_mxfp4 += 1
            elif recipe_tensor.ggml_type == GGML_Q2_K:
                if artifact_recipe.q2_k_error_limit is None:
                    raise ValueError(
                        f"{artifact_recipe.name}: Q2_K tensor has no error limit"
                    )
                q2_k_tensors += 1
                max_q2_k_error = max(
                    max_q2_k_error,
                    _verify_q2_k(
                        source,
                        reader,
                        recipe_tensor,
                        artifact_recipe.q2_k_error_limit,
                    ),
                )
            else:
                raise ValueError(f"unsupported GGML type {recipe_tensor.ggml_type}")
        _verify_zero_padding(reader)
        expected_file_size = reader.data_start + align_up(
            plan[-1].byte_len + reader.tensors[-1].relative_offset
        )
        actual_file_size = gguf_path.stat().st_size
        if expected_file_size != artifact_recipe.expected_file_size:
            raise ValueError(
                f"{artifact_recipe.name}: recipe computes file size "
                f"{expected_file_size}, expected {artifact_recipe.expected_file_size}"
            )
        if actual_file_size != expected_file_size:
            raise ValueError(
                f"GGUF size mismatch: expected {expected_file_size}, found {actual_file_size}"
            )
        data_start = reader.data_start

    all_type_counts = {
        GGML_F32: exact_f32,
        GGML_F16: exact_f16,
        GGML_Q8_0: q8_tensors,
        GGML_Q2_K: q2_k_tensors,
        GGML_MXFP4: exact_mxfp4,
    }
    observed_type_counts = tuple(
        (ggml_type, all_type_counts[ggml_type])
        for ggml_type, _count in artifact_recipe.expected_type_counts
    )
    unexpected_types = {
        ggml_type: count
        for ggml_type, count in all_type_counts.items()
        if count and ggml_type not in dict(artifact_recipe.expected_type_counts)
    }
    if observed_type_counts != artifact_recipe.expected_type_counts or unexpected_types:
        raise ValueError(
            f"unexpected tensor type inventory for {artifact_recipe.name}: "
            f"observed={observed_type_counts}, unexpected={unexpected_types}, "
            f"expected={artifact_recipe.expected_type_counts}"
        )

    type_report = {
        GGML_TYPE_NAMES[ggml_type]: count
        for ggml_type, count in artifact_recipe.expected_type_counts
    }
    validation: dict[str, object] = {
        "f32_exact_tensors": exact_f32,
        "f16_exact_tensors": exact_f16,
        "q8_0_max_block_relative_error": max_q8_error,
        "q8_0_limit": 0.01,
        "status": "pass",
    }
    if exact_mxfp4:
        validation["mxfp4_exact_tensors"] = exact_mxfp4
    if q2_k_tensors:
        validation.update(
            {
                "q2_k_tensors": q2_k_tensors,
                "q2_k_max_block_relative_error": max_q2_k_error,
                "q2_k_limit": artifact_recipe.q2_k_error_limit,
            }
        )

    report: dict[str, object] = {
        "artifact": {
            "filename": gguf_path.name,
            "size": gguf_path.stat().st_size,
            "sha256": sha256_path(gguf_path),
        },
        "source": {
            "repository": SOURCE_REPOSITORY,
            "revision": SOURCE_REVISION,
            "files": sources,
        },
        "gguf": {
            "version": 3,
            "data_start": data_start,
            "alignment": ALIGNMENT,
            "tensor_count": len(plan),
            "types": type_report,
        },
        "recipe": artifact_recipe.name,
        "validation": validation,
    }
    print(json.dumps(report, indent=2, sort_keys=True), flush=True)
    return report


def write_report(path: Path, report: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.partial-{os.getpid()}")
    with temporary.open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(report, handle, indent=2, sort_keys=True)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sources", type=Path, required=True)
    parser.add_argument(
        "--recipe",
        choices=sorted(RECIPES),
        default=DEFAULT_RECIPE.name,
        help=f"Artifact recipe (default: {DEFAULT_RECIPE.name})",
    )
    parser.add_argument(
        "--gguf",
        type=Path,
        help="GGUF path (default: the selected recipe's canonical filename)",
    )
    parser.add_argument("--report", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    recipe = resolve_recipe(args.recipe)
    gguf = args.gguf if args.gguf is not None else Path(recipe.output_filename)
    report = verify(args.sources, gguf, recipe=recipe)
    if args.report:
        write_report(args.report, report)


if __name__ == "__main__":
    main()