"""Capture the exact source and input identity of a calibration run."""
import hashlib
import json
from pathlib import Path
import torch


def capture_sources(output: Path, args) -> dict:
    source_root = Path(__file__).resolve().parents[1]
    destination = output.with_suffix(".sources")
    resuming = getattr(args, "resume", False)
    if not resuming:
        destination.mkdir(parents=True, exist_ok=False)
    hashes = {}
    for path in sorted(source_root.rglob("*")):
        if path.suffix not in {".py", ".yaml"} or "__pycache__" in path.parts:
            continue
        relative = path.relative_to(source_root)
        data = path.read_bytes()
        hashes[str(relative)] = hashlib.sha256(data).hexdigest()
        target = destination / relative
        if not resuming:
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_bytes(data)
    metadata = {
        "source_sha256": hashes, "source_snapshot": str(destination.resolve()),
        "arguments": {k: str(v) if isinstance(v, Path) else v for k,v in vars(args).items()},
        "torch": str(torch.__version__), "cuda": torch.version.cuda,
        "physical_gpu": args.gpu, "gpu_count": 1,
    }
    metadata = json.loads(json.dumps(metadata))
    if resuming:
        previous = json.loads((destination / "manifest.json").read_text())
        ignored = {"resume", "gpu", "stop_after_steps"}
        requested = {key: value for key, value in metadata["arguments"].items() if key not in ignored}
        original = {key: value for key, value in previous["arguments"].items() if key not in ignored}
        # JSON normalizes factor tuples to lists.
        if json.loads(json.dumps(requested)) != original or previous["source_sha256"] != hashes:
            raise ValueError("resume requires unchanged calibration source and arguments")
        return previous
    (destination / "manifest.json").write_text(json.dumps(metadata, indent=2) + "\n")
    return metadata


def text_hashes(texts):
    return [hashlib.sha256(text.encode("utf-8")).hexdigest() for text in texts]