from __future__ import annotations

import hashlib
import json
from collections import Counter
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
from onnxruntime.transformers.float16 import convert_float_to_float16
from onnxruntime.quantization.onnx_model import ONNXModel


ROOT = Path(__file__).resolve().parent
OUTPUT_DIR = ROOT / "output"
REPORT_DIR = ROOT / "reports"
GRAPH_NAMES = ("decoder_prefill", "decoder_step")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def make_inputs(session: ort.InferenceSession) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(20260716)
    feeds: dict[str, np.ndarray] = {}
    for value in session.get_inputs():
        shape = [257 if dimension == "past_len" else dimension for dimension in value.shape]
        if value.name == "vision_embeds":
            feeds[value.name] = rng.normal(0, 0.2, shape).astype(np.float32)
        elif value.name == "token_one_hot":
            token = np.zeros(shape, dtype=np.float32)
            token[0, 0, 4] = 1
            feeds[value.name] = token
        elif value.name == "position_ids":
            feeds[value.name] = np.array([[257]], dtype=np.int32)
        else:
            feeds[value.name] = rng.normal(0, 0.02, shape).astype(np.float32)
    return feeds


def compare(reference: Path, candidate: Path) -> dict:
    expected_session = ort.InferenceSession(
        str(reference), providers=["CPUExecutionProvider"]
    )
    actual_session = ort.InferenceSession(
        str(candidate), providers=["CPUExecutionProvider"]
    )
    feeds = make_inputs(expected_session)
    expected = expected_session.run(None, feeds)
    actual = actual_session.run(None, feeds)
    differences = [
        float(np.max(np.abs(expected_value - actual_value)))
        for expected_value, actual_value in zip(expected, actual)
    ]
    expected_token = int(expected[0][0, -1].argmax())
    actual_token = int(actual[0][0, -1].argmax())
    return {
        "max_abs": max(differences),
        "logits_max_abs": differences[0],
        "cache_max_abs": max(differences[1:]),
        "reference_top_token": expected_token,
        "candidate_top_token": actual_token,
        "top_token_matches": expected_token == actual_token,
    }


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    report = {
        "description": (
            "Complete Baberu decoder converted to FP16 internal tensors with FP32 "
            "public inputs and outputs; no layer, hidden-size, vocabulary, or "
            "generation-logic reduction"
        ),
        "graphs": {},
    }
    for name in GRAPH_NAMES:
        source = OUTPUT_DIR / f"{name}_fp32.onnx"
        destination = OUTPUT_DIR / f"{name}_fp16.onnx"
        if not source.exists():
            raise SystemExit(f"Missing {source}. Run export_decoder_fp32.py first.")
        print(f"Converting {source.name} -> {destination.name}", flush=True)
        model = convert_float_to_float16(onnx.load(source), keep_io_types=True)
        sortable = ONNXModel(model)
        sortable.topological_sort()
        model = sortable.model
        model.producer_name = "vibe-manga-baberu-webgpu-fp16"
        model.producer_version = "1"
        onnx.checker.check_model(model)
        onnx.save(model, destination)
        operators = Counter(node.op_type for node in model.graph.node)
        parity = compare(source, destination)
        if not parity["top_token_matches"]:
            raise RuntimeError(f"{destination.name}: synthetic top-token parity failed")
        report["graphs"][name] = {
            "source": source.name,
            "destination": destination.name,
            "bytes": destination.stat().st_size,
            "sha256": sha256(destination),
            "operators": dict(sorted(operators.items())),
            "fp32_cpu_parity": parity,
        }

    destination = REPORT_DIR / "decoder-fp16-report.json"
    destination.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()