"""Offline assembly of a GAP release from validated run artifacts."""

from __future__ import annotations

import copy
import hashlib
import json
from pathlib import Path
from typing import Any

from .offline import load_dataset
from .store import safe_name, sha256_payload
from .surface import SURFACE_FAMILIES


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 read_stage(run_root: Path, item_id: str, stage: str) -> dict[str, Any]:
    path = (
        run_root
        / "items"
        / safe_name(item_id)
        / "stages"
        / f"{safe_name(stage)}.json"
    )
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload.get("payload_sha256") != sha256_payload(payload.get("payload")):
        raise ValueError(f"stage digest mismatch: {path}")
    return payload["payload"]


def read_kernel(run_root: Path, item_id: str) -> tuple[dict[str, Any], Path]:
    path = run_root / "items" / safe_name(item_id) / "final.json"
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload.get("status") != "accepted" or not payload.get("accepted_candidate"):
        raise ValueError(f"kernel run is not accepted: {path}")
    candidate = payload["accepted_candidate"]
    expected = payload.get("accepted_candidate_sha256")
    if expected and expected != sha256_payload(candidate):
        raise ValueError(f"accepted candidate digest mismatch: {path}")
    return payload, path


def export_release(
    *,
    source_dataset: Path,
    surface_run_root: Path,
    kernel_run_root: Path,
    output_root: Path,
    allow_partial: bool = False,
    item_ids: set[str] | None = None,
) -> dict[str, Any]:
    records_dir = output_root / "records"
    if output_root.exists() and any(output_root.iterdir()):
        raise FileExistsError(f"refusing to overwrite non-empty {output_root}")
    records_dir.mkdir(parents=True, exist_ok=True)

    source = load_dataset(source_dataset)
    if item_ids is not None:
        unknown = item_ids - set(source)
        if unknown:
            raise ValueError(f"unknown requested item IDs: {sorted(unknown)}")
        source = {
            item_id: record
            for item_id, record in source.items()
            if item_id in item_ids
        }
    exported: list[dict[str, Any]] = []
    missing: list[dict[str, str]] = []
    for item_id, record in sorted(source.items()):
        try:
            variants = {
                family: read_stage(
                    surface_run_root,
                    item_id,
                    f"surface_{family}_variant",
                )
                for family in SURFACE_FAMILIES
            }
            kernel_payload, kernel_path = read_kernel(kernel_run_root, item_id)
        except (FileNotFoundError, ValueError) as exc:
            missing.append({"item_id": item_id, "reason": str(exc)})
            continue

        candidate = kernel_payload["accepted_candidate"]
        variants["kernel_variant"] = {
            "question": candidate["question"],
            "solution": candidate["solution"],
        }
        output_record = copy.deepcopy(record)
        output_record["variants"] = variants
        output_path = records_dir / f"{safe_name(item_id)}.json"
        output_path.write_text(
            json.dumps(output_record, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        exported.append(
            {
                "item_id": item_id,
                "path": str(output_path.resolve()),
                "sha256": sha256_file(output_path),
            }
        )

    if missing and not allow_partial:
        # Keep the manifest for diagnosis, but never leave a partial directory
        # that looks like a completed release.
        status = "incomplete"
    else:
        status = "complete" if not missing else "partial_allowed"
    manifest = {
        "status": status,
        "source_dataset": str(source_dataset.resolve()),
        "surface_run_root": str(surface_run_root.resolve()),
        "kernel_run_root": str(kernel_run_root.resolve()),
        "source_item_count": len(source),
        "exported_item_count": len(exported),
        "missing_item_count": len(missing),
        "allow_partial": allow_partial,
        "exported": exported,
        "missing": missing,
    }
    (output_root / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    if missing and not allow_partial:
        raise RuntimeError(
            f"release incomplete: exported {len(exported)}/{len(source)}; "
            f"see {output_root / 'manifest.json'}"
        )
    return manifest