"""Build datasets, splits, manifests. Agent-visible documents go to data/generated/docs/, answer keys and hidden tests
to data/generated/keys/ (never mounted into prompts by the harness; checked by tests/test_leakage.py)."""
from __future__ import annotations
import json, yaml, hashlib
from pathlib import Path
from ..util import ROOT, write_json, read_json, sha256_text, utc_now
from .synthetic import generate_document
from .exec_tasks import generate_exec_task
from .public import load_subset

GEN = ROOT / "data" / "generated"
DOCS, KEYS, MAN = GEN / "docs", GEN / "keys", ROOT / "data" / "manifests"

def _write_item(item: dict, family: str, split: str, extra_key: dict | None = None):
    did = item["document_id"]
    doc_pub = {"document_id": did, "family": family, "split": split, "text": item["text"], "char_len": item["char_len"],
               "document_hash": sha256_text(item["text"]),
               "questions": [{k: q[k] for k in ("question_id", "qtype", "question")} for q in item["questions"]]}
    key = {"document_id": did, "family": family, "split": split, "document_hash": doc_pub["document_hash"],
           "questions": item["questions"], "facts": item.get("facts", [])}
    if extra_key:
        key.update(extra_key)
    (DOCS / family).mkdir(parents=True, exist_ok=True); (KEYS / family).mkdir(parents=True, exist_ok=True)
    write_json(DOCS / family / f"{did}.json", doc_pub)
    write_json(KEYS / family / f"{did}.json", key)
    return {"document_id": did, "family": family, "split": split, "document_hash": doc_pub["document_hash"], "char_len": item["char_len"],
            "n_questions": len(item["questions"]), "cluster_id": did}

def prepare(config_path: str) -> int:
    cfg = yaml.safe_load(open(config_path))
    manifest = {"created_utc": utc_now(), "config": cfg, "config_hash": sha256_text(json.dumps(cfg, sort_keys=True)), "items": []}
    # synthetic
    sc = cfg["synthetic"]; L = sc["lengths_tokens"]
    for split, r in sc["splits"].items():
        for i in range(r["count"]):
            seed = r["start"] + i; tl = L[i % len(L)]
            d = generate_document(seed, sc["n_projects"], tl)
            manifest["items"].append(_write_item(d, "synthetic", split) | {"target_tokens": tl, "seed": seed})
    ec = cfg["exec"]
    for split, r in ec["splits"].items():
        for i in range(r["count"]):
            seed = r["start"] + i
            t = generate_exec_task(seed, ec["target_tokens"])
            manifest["items"].append(_write_item(t, "exec", split, {"hidden_tests": t["hidden_tests"]}) | {"seed": seed})
    for sub in cfg["public"]["subsets"] + cfg["public"].get("extension_subsets", []):
        for it in load_subset(sub):
            h = it["document_hash"]
            split = "dev" if int(h[-1], 16) <= 3 else "test"
            fam = f"public_{sub}"
            manifest["items"].append(_write_item(it, fam, split) | {"length_words": it["length_words"], "extension": sub in cfg["public"].get("extension_subsets", [])})
    # duplicate check across splits (exact hash)
    seen = {}
    dups = []
    for it in manifest["items"]:
        if it["document_hash"] in seen and seen[it["document_hash"]] != it["split"]:
            dups.append(it["document_id"])
        seen[it["document_hash"]] = it["split"]
    manifest["cross_split_exact_duplicates"] = dups
    MAN.mkdir(parents=True, exist_ok=True)
    write_json(MAN / "splits.json", manifest)
    # license/provenance manifest
    write_json(MAN / "datasets.json", {
        "synthetic": {"source": "wp30.data.synthetic (this repository)", "license": "CC BY 4.0 (generated)", "note": "Fictional entities; a procedurally generated ledger is not a real filing."},
        "exec": {"source": "wp30.data.exec_tasks (this repository)", "license": "MIT (generated)"},
        "public_longbench": {"source": "https://huggingface.co/datasets/THUDM/LongBench data.zip", "paper": "Bai et al. 2024, ACL (10.18653/v1/2024.acl-long.172)",
                             "subsets": {"qasper": "CC BY 4.0 (upstream)", "2wikimqa": "Apache-2.0 repo (data licence unverified)", "hotpotqa": "CC BY-SA 4.0 (upstream)"},
                             "official_metric": "F1 (LongBench normalization)", "note": "LongBench dataset card has no licence field (unverified)."},
    })
    from collections import Counter
    c = Counter((it["family"], it["split"]) for it in manifest["items"])
    for k in sorted(c): print(k, c[k])
    print("cross-split exact duplicates:", len(dups))
    return 0

def validate(manifest_path: str) -> int:
    man = read_json(manifest_path)
    ok = True
    for it in man["items"]:
        p = DOCS / it["family"] / f"{it['document_id']}.json"
        if not p.exists():
            print("missing", p); ok = False; continue
        d = read_json(p)
        if sha256_text(d["text"]) != it["document_hash"]:
            print("hash mismatch", it["document_id"]); ok = False
        for q in d["questions"]:
            if "answer" in q or "answers" in q:
                print("LEAK: answer in public doc", it["document_id"]); ok = False
        if "facts" in d or "hidden_tests" in d:
            print("LEAK: facts/tests in public doc", it["document_id"]); ok = False
    print("validate:", "OK" if ok else "FAILED", len(man["items"]), "items")
    return 0 if ok else 1

def load_doc(document_id: str, family: str) -> dict:
    return read_json(DOCS / family / f"{document_id}.json")

def load_key(document_id: str, family: str) -> dict:
    return read_json(KEYS / family / f"{document_id}.json")