#!/usr/bin/env python3
"""Publish checksums for the completed AD01 compiled numerical package."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import tempfile
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", required=True, type=Path)
    parser.add_argument("--result-dir", required=True, type=Path)
    parser.add_argument("--report", action="append", default=[], type=Path)
    return parser.parse_args()


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


def atomic_text(path: Path, value: str) -> None:
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        handle.write(value)
        temporary = Path(handle.name)
    os.replace(temporary, path)


def main() -> int:
    args = parse_args()
    root = args.repo_root.resolve()
    results = args.result_dir.resolve()
    validation = results / "validation.json"
    if not validation.is_file():
        raise ValueError("terminal validation.json is required before manifest publication")
    validation_status = json.loads(validation.read_text()).get("status")
    if validation_status not in {"PASS", "FAIL"}:
        raise ValueError("validation.json must have a terminal PASS or FAIL status")
    merge_summary = results / "compiled_merge_summary.json"
    if not merge_summary.is_file():
        raise ValueError("compiled_merge_summary.json is required")
    fidelity_status = json.loads(merge_summary.read_text()).get("fidelity_status")
    excluded = {results / "artifact_manifest.json", results / "artifacts.sha256"}
    temporary_files = sorted(
        path
        for path in results.rglob("*")
        if path.is_file()
        and (
            path.name.startswith("tmp")
            or path.name.endswith(".tmp")
        )
    )
    if temporary_files:
        raise ValueError(
            "temporary files must be removed before publication: "
            + ", ".join(str(path) for path in temporary_files)
        )
    files = sorted(
        path for path in results.rglob("*") if path.is_file() and path not in excluded
    )
    external_reports = [path.resolve() for path in args.report]
    for path in external_reports:
        if not path.is_file():
            raise FileNotFoundError(path)
    records = []
    for path in files + external_reports:
        try:
            display = str(path.relative_to(root))
        except ValueError:
            display = str(path)
        records.append(
            {
                "path": display,
                "bytes": path.stat().st_size,
                "sha256": sha256_file(path),
                "scope": "RESULT_PACKAGE" if path in files else "EXTERNAL_REPORT",
            }
        )
    manifest = {
        "schema_version": "1.0",
        "model_id": "AD01",
        "stage": "S6-NV_COMPILED_MLIR_NUMERICAL_AND_Q1_VALIDATION",
        "status": (
            "PASS"
            if validation_status == "PASS" and fidelity_status == "PASS"
            else "PARTIAL"
        ),
        "independent_validation_status": validation_status,
        "numerical_fidelity_status": fidelity_status,
        "files": records,
        "policy": {
            "latency_measured": False,
            "converter_run": False,
            "lowering_run": False,
            "codegen_run": False,
            "training_or_calibration": False,
            "model_or_weight_modified": False,
        },
    }
    manifest_path = results / "artifact_manifest.json"
    atomic_text(manifest_path, json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    checksum_records = records + [
        {
            "path": str(manifest_path.relative_to(root)),
            "sha256": sha256_file(manifest_path),
        }
    ]
    atomic_text(
        results / "artifacts.sha256",
        "".join(f"{row['sha256']}  {row['path']}\n" for row in checksum_records),
    )
    print(json.dumps({"status": "PASS", "files": len(records), "manifest": str(manifest_path)}, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())