"""Verify an explicit training-data expansion without changing evaluation data."""
import hashlib
import json
from pathlib import Path


SPLITS = ("train", "validation", "calibration", "test", "holdout")
PROTECTED_SPLITS = SPLITS[1:]


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def verified_manifest(dataset):
    root = Path(dataset).resolve()
    manifest = json.loads((root / "manifest.json").read_text())
    hashes = manifest.get("split_sha256", {})
    if set(hashes) != set(SPLITS):
        raise ValueError("Data transition requires exactly the five registered splits")
    for split, expected in hashes.items():
        path = root / f"{split}.jsonl"
        if not path.is_file() or sha256(path) != expected:
            raise ValueError(f"Data transition checksum mismatch: {root}/{split}")
    return manifest


def data_signature(manifest, provenance, max_tokens):
    return hashlib.sha256(json.dumps({
        "data": manifest["split_sha256"], "model": provenance,
        "implementation": sha256(Path(__file__).with_name("training_model.py")),
        "max_tokens": max_tokens,
    }, sort_keys=True).encode()).hexdigest()


def verify_train_data_transition(saved, config, signature):
    """Check both real datasets and signatures; return immutable lineage evidence."""
    parent_config = saved["config"]
    if parent_config.get("max_tokens") != config["max_tokens"]:
        raise ValueError("Training-data expansion must preserve max_tokens")
    parent_root = Path(parent_config["dataset"]).resolve()
    current_root = Path(config["dataset"]).resolve()
    parent = verified_manifest(parent_root)
    current = verified_manifest(current_root)
    provenance = saved["model_provenance"]
    if data_signature(parent, provenance, config["max_tokens"]) != saved["data_signature"]:
        raise ValueError("Parent data signature does not match its verified source")
    if data_signature(current, provenance, config["max_tokens"]) != signature:
        raise ValueError("Expanded data signature does not match its verified source")
    protected = {split: parent["split_sha256"][split] for split in PROTECTED_SPLITS}
    if any(current["split_sha256"][split] != digest for split, digest in protected.items()):
        raise ValueError("Training-data expansion must preserve every reserved evaluation split")
    base = current.get("base_dataset", {})
    if (Path(base.get("path", "")).resolve() != parent_root or
            base.get("manifest_sha256") != sha256(parent_root / "manifest.json") or
            base.get("split_sha256") != parent["split_sha256"]):
        raise ValueError("Expanded data must record the verified parent dataset lineage")
    if current["split_sha256"]["train"] == parent["split_sha256"]["train"]:
        raise ValueError("Training-data expansion did not change the training split")
    return {"kind": "training_split_only", "parent_dataset": str(parent_root),
            "dataset": str(current_root), "parent_data_signature": saved["data_signature"],
            "data_signature": signature, "parent_manifest_sha256": sha256(parent_root / "manifest.json"),
            "manifest_sha256": sha256(current_root / "manifest.json"),
            "protected_split_sha256": protected}