#!/usr/bin/env python3
"""Compare deterministic evidence from two DeepSeek-V2-Lite trace runs."""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--first", type=Path, required=True)
    parser.add_argument("--second", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--captured-at", default=None)
    return parser.parse_args()


def read(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    args = parse_args()
    first = read(args.first)
    second = read(args.second)
    checks: list[dict[str, Any]] = []

    def check(name: str, left: Any, right: Any) -> None:
        checks.append({"name": name, "exact": left == right})

    check("provenance", first["provenance"], second["provenance"])
    check("configuration", first["configuration"], second["configuration"])
    check("prompt_tokenization", first["prompts"], second["prompts"])
    check(
        "initial_hidden_hash",
        first["execution"]["initial_hidden"]["sha256_fp32"],
        second["execution"]["initial_hidden"]["sha256_fp32"],
    )
    check(
        "final_hidden_hash",
        first["execution"]["final_hidden"]["sha256_fp32"],
        second["execution"]["final_hidden"]["sha256_fp32"],
    )
    for left, right in zip(
        first["execution"]["layers"],
        second["execution"]["layers"],
        strict=True,
    ):
        layer = left["layer"]
        check(
            f"layer_{layer}_hidden_hash",
            left["hidden_after"]["sha256_fp32"],
            right["hidden_after"]["sha256_fp32"],
        )
        check(f"layer_{layer}_mla_shapes", left["mla"], right["mla"])
        if "routing" in left:
            check(
                f"layer_{layer}_aggregate_load",
                left["routing"]["aggregate_load"],
                right["routing"]["aggregate_load"],
            )
            check(
                f"layer_{layer}_token_routes",
                [
                    prompt["token_routes"]
                    for prompt in left["routing"]["per_prompt"]
                ],
                [
                    prompt["token_routes"]
                    for prompt in right["routing"]["per_prompt"]
                ],
            )

    captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
    result = {
        "schema_version": 1,
        "captured_at": captured_at,
        "first": {
            "captured_at": first["captured_at"],
            "sha256": sha256(args.first),
        },
        "second": {
            "captured_at": second["captured_at"],
            "sha256": sha256(args.second),
        },
        "timing_compared": False,
        "checks": checks,
        "exact_checks": sum(item["exact"] for item in checks),
        "total_checks": len(checks),
        "all_exact": all(item["exact"] for item in checks),
    }
    if not result["all_exact"]:
        failed = [item["name"] for item in checks if not item["exact"]]
        raise AssertionError(f"trace mismatch: {failed}")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()