"""Pinned, checksum-verified data preparation without evidence truncation."""

import json
from pathlib import Path

from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

from .codec import Codec
from .common import ROOT, SPLITS, atomic_json, digest, sha_file


def tokenizer_for(cfg, directory=None):
    tokenizer = AutoTokenizer.from_pretrained(
        str(directory or cfg["model"]["repo"]),
        **({} if directory else {"revision": cfg["model"]["revision"]}),
    )
    return Codec(tokenizer, cfg["model"]["max_length"], cfg["model"]["max_options"])


def prepared_directory(cfg):
    path = Path(cfg["data"]["directory"])
    return path if path.is_absolute() else ROOT / path


def subset(record, names):
    return {
        "input": {
            "state": record["input"]["state"],
            "questions": {k: record["input"]["questions"][k] for k in names},
        },
        **{
            field: {k: record[field][k] for k in names}
            for field in ("targets", "option_keys", "decision_ids", "categories")
        },
    }


def prepare(cfg, output=None):
    out = Path(output) if output else prepared_directory(cfg)
    if (out / "manifest.json").exists():
        verify_prepared(cfg, out)
        return out
    out.mkdir(parents=True, exist_ok=True)
    codec = tokenizer_for(cfg)
    repo, revision = cfg["data"]["repo"], cfg["data"]["revision"]

    def download(name):
        return Path(hf_hub_download(repo, name, repo_type="dataset", revision=revision))

    checks = json.loads(download("CHECKSUMS.json").read_text())

    def verified(name):
        path = download(name)
        if sha_file(path) != checks[name]["sha256"]:
            raise ValueError(f"Pinned dataset checksum mismatch: {name}")
        return path

    result = {
        "dataset": {"repo": repo, "revision": revision},
        "model": cfg["model"],
        "splits": {},
        "codec": codec.config(),
    }
    all_ids = set()
    for split in SPLITS:
        metadata = {}
        with verified(f"original/decisions/{split}.jsonl").open() as stream:
            for line in stream:
                row = json.loads(line)
                metadata[row["id"]] = row["metadata"]["category"]
        seen, count, split_requests, max_tokens = set(), 0, 0, 0
        temporary = out / (split + ".jsonl.tmp")
        with (
            verified(f"original/requests/{split}.jsonl").open() as source,
            temporary.open("w") as dest,
        ):
            for line in source:
                record = json.loads(line)
                record["categories"] = {k: metadata[v] for k, v in record["decision_ids"].items()}
                packs, names = [], []
                for name in record["input"]["questions"]:
                    attempt = subset(record, names + [name])
                    try:
                        codec.encode(
                            attempt["input"]["state"],
                            attempt["input"]["questions"],
                            attempt["targets"],
                            attempt["option_keys"],
                        )
                        names.append(name)
                    except ValueError as exc:
                        if "Rendered request has" not in str(exc) or not names:
                            raise
                        packs.append(subset(record, names))
                        names = [name]
                        single = subset(record, names)
                        codec.encode(
                            single["input"]["state"],
                            single["input"]["questions"],
                            single["targets"],
                            single["option_keys"],
                        )
                if names:
                    packs.append(subset(record, names))
                for pack in packs:
                    ids = list(pack["decision_ids"].values())
                    if seen.intersection(ids) or all_ids.intersection(ids):
                        raise ValueError("Duplicated or cross-split decision")
                    seen.update(ids)
                    encoded = codec.encode(
                        pack["input"]["state"],
                        pack["input"]["questions"],
                        pack["targets"],
                        pack["option_keys"],
                    )
                    pack["encoded"] = encoded
                    pack["id"] = digest([split, ids])
                    dest.write(json.dumps(pack, ensure_ascii=False, separators=(",", ":")) + "\n")
                    count += len(ids)
                    split_requests += 1
                    max_tokens = max(max_tokens, len(encoded["input_ids"]))
        if seen != set(metadata):
            raise ValueError(f"Lost or invented decisions in {split}")
        all_ids.update(seen)
        target = out / (split + ".jsonl")
        temporary.replace(target)
        result["splits"][split] = {
            "decisions": count,
            "requests": split_requests,
            "max_tokens": max_tokens,
            "sha256": sha_file(target),
        }
        print(f"Prepared {split}: {count:,} decisions in {split_requests:,} requests", flush=True)
    codec.tokenizer.save_pretrained(out / "tokenizer")
    result["config_hash"] = digest(
        {"model": cfg["model"], "data": {"repo": repo, "revision": revision}}
    )
    result["manifest_hash"] = digest(result)
    atomic_json(out / "manifest.json", result)
    return out


def verify_prepared(cfg, directory=None):
    directory = Path(directory) if directory else prepared_directory(cfg)
    manifest = json.loads((directory / "manifest.json").read_text())
    expected = digest(
        {
            "model": cfg["model"],
            "data": {"repo": cfg["data"]["repo"], "revision": cfg["data"]["revision"]},
        }
    )
    if manifest["config_hash"] != expected:
        raise ValueError("Prepared data belongs to another model/data configuration")
    check = dict(manifest)
    stored = check.pop("manifest_hash")
    if digest(check) != stored:
        raise ValueError("Prepared manifest was modified")
    for split in SPLITS:
        if sha_file(directory / (split + ".jsonl")) != manifest["splits"][split]["sha256"]:
            raise ValueError(f"Prepared checksum mismatch: {split}")
    return manifest


def load_rows(directory, split):
    with (Path(directory) / (split + ".jsonl")).open() as stream:
        return [json.loads(line) for line in stream]


def encoded_row(row, codec, seed=None):
    if seed is None:
        return row["encoded"]
    return codec.encode(
        row["input"]["state"],
        row["input"]["questions"],
        row["targets"],
        row["option_keys"],
        seed=seed,
    )