"""Explicit, audited continuation of the existing run to three epochs."""

import json
import shutil
from pathlib import Path

from .checkpoint import checkpoint_path
from .common import atomic_json, digest, sha_file, source_hash

LIMITS = {"budget_usd": 15.0, "stop_cost_usd": 14.0, "max_worker_seconds": 18000}


def make_continuation(cfg, root):
    root = Path(root)
    checkpoint = checkpoint_path(root)
    return {
        "kind": "finish_three_epochs",
        "parent_identity": json.loads((root / "run.json").read_text()),
        "accepted_code_hash": source_hash(),
        "config_hash": digest(cfg),
        "checkpoint": str(checkpoint.relative_to(root)),
        "checkpoint_training_sha256": sha_file(checkpoint / "training.pt"),
        "limits": LIMITS,
    }


def validate_continuation(cfg, root, continuation):
    root = Path(root)
    identity = json.loads((root / "run.json").read_text())
    if (
        continuation.get("kind") != "finish_three_epochs"
        or continuation.get("limits") != LIMITS
        or cfg["train"]["epochs"] != 3
        or continuation.get("parent_identity") != identity
        or continuation.get("accepted_code_hash") != source_hash()
        or continuation.get("config_hash") != digest(cfg)
        or identity["config_hash"] != digest(cfg)
    ):
        raise ValueError("Continuation identity, code, configuration or authorized limits changed")
    checkpoint = checkpoint_path(root)
    if str(checkpoint.relative_to(root)) != continuation.get("checkpoint") or sha_file(
        checkpoint / "training.pt"
    ) != continuation.get("checkpoint_training_sha256"):
        raise ValueError("Continuation must use the verified latest checkpoint")
    return checkpoint


def archive_previous_export(root):
    root = Path(root)
    archive = root / "archives" / "before-three-epoch-continuation"
    if archive.exists():
        return
    archive.mkdir(parents=True)
    for name in (
        "export",
        "summary.json",
        "evaluation.json",
        "calibration.json",
        "test-predictions.jsonl",
        "local-reload-verification.json",
        "budget.json",
        "best.json",
        "latest.json",
        "provenance.json",
    ):
        source = root / name
        if source.is_dir():
            shutil.copytree(source, archive / name)
        elif source.exists():
            shutil.copy2(source, archive / name)
    atomic_json(archive / "complete.json", {"archived": True})