#!/usr/bin/env python3
"""Link a frozen L3 run to its pre-training development evaluation plan."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path


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 load(path: Path, artifact_type: str) -> dict:
    value = json.loads(path.read_text(encoding="utf-8"))
    if value.get("artifact_type") != artifact_type:
        raise ValueError(f"{path} has the wrong artifact_type")
    return value


def resolve_recorded(value: str, receipt_path: Path) -> Path:
    path = Path(value)
    if not path.is_absolute():
        path = receipt_path.parent / path
    return path.resolve()


def portable(path: Path, base: Path) -> str:
    try:
        return str(path.resolve().relative_to(base.resolve()))
    except ValueError:
        return str(path.resolve())


def prepare(args: argparse.Namespace) -> dict:
    build_receipt = load(
        args.bundle_build_receipt, "telen_l3_development_bundle_build"
    )
    training_receipt = load(
        args.training_receipt, "telen_l3_development_pilot"
    )
    if build_receipt.get("authorized") is not True:
        raise ValueError("bundle build receipt is not authorized")
    if training_receipt.get("status") != "completed_development_only":
        raise ValueError("training receipt is incomplete")
    if training_receipt.get("final_test_accessed") is not False:
        raise ValueError("training receipt does not preserve final-test sealing")

    training_manifest = resolve_recorded(
        str(build_receipt["training_manifest"]), args.bundle_build_receipt
    )
    evaluation_plan_path = resolve_recorded(
        str(build_receipt["evaluation_plan"]), args.bundle_build_receipt
    )
    development_queries = resolve_recorded(
        str(build_receipt["development_queries"]), args.bundle_build_receipt
    )
    development_candidates = resolve_recorded(
        str(build_receipt["development_candidates"]), args.bundle_build_receipt
    )
    expected_hashes = {
        training_manifest: build_receipt["training_manifest_sha256"],
        evaluation_plan_path: build_receipt["evaluation_plan_sha256"],
        development_queries: build_receipt["development_queries_sha256"],
        development_candidates: build_receipt["development_candidates_sha256"],
    }
    for path, expected in expected_hashes.items():
        if not path.is_file() or sha256_file(path) != expected:
            raise ValueError(f"frozen build artifact changed: {path}")
    if training_receipt.get("manifest_sha256") != sha256_file(training_manifest):
        raise ValueError("training receipt belongs to another training manifest")
    evaluation_plan = load(
        evaluation_plan_path, "telen_l3_development_evaluation_plan"
    )
    if evaluation_plan.get("final_test_accessed") is not False:
        raise ValueError("evaluation plan does not preserve final-test sealing")
    if (
        evaluation_plan.get("development_queries_sha256")
        != sha256_file(development_queries)
        or evaluation_plan.get("development_candidates_sha256")
        != sha256_file(development_candidates)
    ):
        raise ValueError("evaluation plan input hashes changed")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    files = {
        "training_manifest": training_manifest,
        "training_receipt": args.training_receipt,
        "development_queries": development_queries,
        "development_candidates": development_candidates,
    }
    manifest = {
        "schema_version": 1,
        "artifact_type": "telen_l3_development_evaluation_bundle",
        "pretraining_evaluation_plan": str(evaluation_plan_path),
        "pretraining_evaluation_plan_sha256": sha256_file(evaluation_plan_path),
        "files": {
            name: {
                "path": portable(path, args.output.parent),
                "sha256": sha256_file(path),
            }
            for name, path in files.items()
        },
        "evaluation": evaluation_plan["evaluation"],
    }
    args.output.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return manifest


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bundle-build-receipt", type=Path, required=True)
    parser.add_argument("--training-receipt", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    print(json.dumps(prepare(parse_args()), ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()