"""Build a deterministic, bounded, auditable first Themis data release."""

import argparse
import json
import math
import os
import sqlite3
from collections import Counter, defaultdict, deque
from pathlib import Path

import pyarrow.parquet as pq
import yaml
from tokenizers import Tokenizer

from .commitpack import CommitPool
from .fetch import ROOT, sha256_file
from .filters import BenchmarkIndex, category, reusable
from .pack import pack_release
from .schema import InvalidRecord, adapt, digest, state_hash, text

SPLITS = {"train": 0, "validation": 1, "calibration": 2, "test": 3}


def rows(root, source, wanted=None):
    for file in source["files"]:
        split = Path(file["path"]).name.split("-")[0]
        if source["adapter"] == "commitpack":
            split = "train"
        if wanted and split not in wanted:
            continue
        path = root / "data/raw" / source["repo"].replace("/", "--") / file["path"]
        if path.suffix == ".jsonl":
            with path.open() as stream:
                for line in stream:
                    yield json.loads(line), split, file["path"]
        else:
            for batch in pq.ParquetFile(path).iter_batches(batch_size=256):
                for row in batch.to_pylist():
                    yield row, split, file["path"]


def fraction(value, seed):
    return int(digest([seed, value])[:12], 16) / 16**12


class Groups:
    def __init__(self):
        self.parent = {}
        self.rank = {}

    def find(self, key):
        self.parent.setdefault(key, key)
        if self.parent[key] != key:
            self.parent[key] = self.find(self.parent[key])
        return self.parent[key]

    def join(self, a, b):
        a, b = self.find(a), self.find(b)
        if a != b:
            # Lexical order keeps the representative independent of traversal order.
            small, large = sorted((a, b))
            self.parent[large] = small


def initial_split(record, seed):
    m = record["metadata"]
    if "derived_split" in m:
        return m["derived_split"]
    upstream = m["upstream_split"]
    draw = fraction(m["group_key"], seed)
    if upstream == "test":
        return "test"
    if upstream == "validation":
        return "validation" if draw < 0.5 else "calibration"
    if "LocalLLaMA" in m["repo"]:
        if draw < 0.05:
            return "validation"
        if draw < 0.10:
            return "calibration"
    return "train"


def candidate_rows(root, lock, cfg, mappings, stats):
    benchmark = next(s for s in lock["sources"] if s["adapter"] == "benchmark")
    index = BenchmarkIndex(row for row, _, _ in rows(root, benchmark))
    stats["benchmark"] = {
        "rows": index.rows,
        "anchors": index.anchors,
        "method": "exact state plus literal normalized 32-word fragments",
    }
    reserved_states, reserved_groups = set(), set()
    for source in lock["sources"]:
        if source["adapter"] == "benchmark":
            continue
        for row, split, _ in rows(root, source, {"validation", "test"}):
            try:
                reserved_states.add(state_hash(text(row["state"])))
            except InvalidRecord:
                stats["invalid_native_reference_states"] = (
                    stats.get("invalid_native_reference_states", 0) + 1
                )
            group = row.get("group_id") or (
                row.get("workflow") if "synthetic-typed" in source["repo"] else row["id"]
            )
            reserved_groups.add(source["repo"] + ":" + str(group))
    stats["reserved_native_states"] = len(reserved_states)
    tokenizer_meta = lock["tokenizer"]
    tok_path = root / "data/raw" / tokenizer_meta["repo"].replace("/", "--") / "tokenizer.json"
    tokenizer = Tokenizer.from_file(str(tok_path))
    tokenizer.no_truncation()
    tokenizer.no_padding()
    budget = cfg["max_input_tokens"] - cfg["prompt_reserve_tokens"]
    rejected = Counter()
    by_source = Counter()
    for source in lock["sources"]:
        if source["adapter"] == "benchmark":
            continue
        print("Scanning " + source["repo"], flush=True)
        commit_pool = None
        if source["adapter"] == "commitpack":
            path = (
                root / "data/raw" / source["repo"].replace("/", "--") / source["files"][0]["path"]
            )
            commit_pool = CommitPool(path, cfg["seed"])
        for row, split, filename in rows(root, source):
            stats["input_cases"] += 1
            if source["adapter"] == "tasksource":
                if row["source"].startswith("procedural-typed-decisions/"):
                    rejected["procedural_duplicate_collection"] += 1
                    continue
                if row.get("variant", "direct") not in ("direct", "packed_derived"):
                    rejected["redundant_augmentation_variant"] += 1
                    continue
                if not any(
                    row["source"].startswith(p)
                    for ps in mappings["tasksource"].values()
                    for p in ps
                ):
                    rejected["unmapped_source"] += 1
                    continue
            if "procedural" in source["repo"] and split == "train":
                if fraction(row["id"], cfg["seed"]) >= cfg["procedural_case_fraction"]:
                    rejected["procedural_case_subsample"] += 1
                    continue
            try:
                decisions = list(
                    commit_pool.adapt(row, source) if commit_pool else adapt(row, source, split)
                )
            except (InvalidRecord, KeyError, TypeError, ValueError) as exc:
                rejected["normalization:" + str(exc)[:80]] += 1
                continue
            for record in decisions:
                m = record["metadata"]
                name, rule = category(record, mappings)
                if name is None:
                    rejected[rule] += 1
                    continue
                if not reusable(record):
                    rejected["reuse_terms_excluded"] += 1
                    continue
                m.update(category=name, category_rule=rule, shard=filename)
                state = record["input"]["state"]
                if split == "train" and (
                    m["state_hash"] in reserved_states or m["group_key"] in reserved_groups
                ):
                    rejected["native_holdout_overlap"] += 1
                    continue
                match = index.match(state)
                if match:
                    rejected["benchmark_overlap"] += 1
                    stats["benchmark_matches"].append(
                        {
                            "id": record["id"],
                            "repo": source["repo"],
                            "source_id": m["source_id"],
                            "match": match,
                        }
                    )
                    continue
                serialized = json.dumps(record["input"], ensure_ascii=False, separators=(",", ":"))
                n_tokens = len(tokenizer.encode(serialized, add_special_tokens=False).ids)
                if n_tokens > budget:
                    rejected["over_token_budget"] += 1
                    continue
                m["payload_tokens"] = n_tokens
                m["split"] = initial_split(record, cfg["seed"])
                m["sampling_priority"] = digest([cfg["seed"], m["fingerprint"]])
                by_source[m["upstream_source"]] += 1
                yield record
        print("Accepted candidates so far: " + str(sum(by_source.values())), flush=True)
    stats["rejections"] = dict(rejected)
    stats["candidate_sources"] = dict(by_source)


def store_candidates(root, lock, cfg, mappings, stats):
    work = root / "data/work"
    work.mkdir(parents=True, exist_ok=True)
    db_path = work / "candidates.sqlite"
    if db_path.exists():
        db_path.unlink()
    db = sqlite3.connect(db_path)
    db.execute("PRAGMA journal_mode=WAL")
    db.execute(
        "CREATE TABLE candidates (fingerprint TEXT PRIMARY KEY, category TEXT, source TEXT, kind TEXT, label INTEGER, priority TEXT, group_key TEXT, state_key TEXT, split TEXT, conflict INTEGER DEFAULT 0, record TEXT)"
    )
    groups = Groups()
    seen_splits = []
    duplicates, conflicts = 0, 0
    for i, record in enumerate(candidate_rows(root, lock, cfg, mappings, stats)):
        m = record["metadata"]
        group, state = "g:" + m["group_key"], "s:" + m["state_hash"]
        groups.join(group, state)
        seen_splits.append((group, m["split"]))
        previous = db.execute(
            "SELECT record FROM candidates WHERE fingerprint=?", (m["fingerprint"],)
        ).fetchone()
        if previous:
            duplicates += 1
            old = json.loads(previous[0])["metadata"]["canonical_target"]
            if any(abs(a - b) > 1e-5 for a, b in zip(old, m["canonical_target"], strict=True)):
                db.execute(
                    "UPDATE candidates SET conflict=1 WHERE fingerprint=?", (m["fingerprint"],)
                )
                conflicts += 1
            continue
        target = record["target"]
        kind = record["input"]["questions"]["decision"]["type"]
        db.execute(
            "INSERT INTO candidates VALUES (?,?,?,?,?,?,?,?,?,0,?)",
            (
                m["fingerprint"],
                m["category"],
                m["upstream_source"],
                kind,
                max(range(len(target)), key=target.__getitem__),
                m["sampling_priority"],
                group,
                state,
                m["split"],
                json.dumps(record, ensure_ascii=False),
            ),
        )
        if i % 2000 == 0:
            db.commit()
    component_splits = {}
    for group, split in seen_splits:
        key = groups.find(group)
        old = component_splits.get(key, "train")
        component_splits[key] = max((old, split), key=SPLITS.__getitem__)
    updates = [
        (component_splits[groups.find(g)], fp)
        for fp, g in db.execute("SELECT fingerprint,group_key FROM candidates")
    ]
    db.executemany("UPDATE candidates SET split=? WHERE fingerprint=?", updates)
    db.execute("CREATE INDEX sample_idx ON candidates(split,category,priority)")
    db.commit()
    stats.update(
        duplicate_decisions=duplicates,
        conflicting_duplicates=conflicts,
        candidate_decisions=db.execute(
            "SELECT COUNT(*) FROM candidates WHERE conflict=0"
        ).fetchone()[0],
    )
    return db


def select(db, split, targets, source_cap):
    """Round-robin sources, then types/labels, within each category."""
    picked = []
    source_counts = Counter()
    for name, target in targets.items():
        buckets = defaultdict(lambda: defaultdict(deque))
        for fp, source, kind, label in db.execute(
            "SELECT fingerprint,source,kind,label FROM candidates WHERE split=? AND category=? AND conflict=0 ORDER BY priority",
            (split, name),
        ):
            buckets[source][(kind, label)].append(fp)
        sources = deque(sorted(buckets))
        strata = {source: deque(sorted(buckets[source])) for source in sources}
        count = 0
        while sources and count < target:
            source = sources.popleft()
            if source_counts[source] >= source_cap:
                continue
            options = strata[source]
            while options and not buckets[source][options[0]]:
                options.popleft()
            if not options:
                continue
            key = options.popleft()
            picked.append(buckets[source][key].popleft())
            source_counts[source] += 1
            count += 1
            if buckets[source][key]:
                options.append(key)
            if options:
                sources.append(source)
    return picked


def export(root, db, cfg, lock, stats):
    out = root / "data/curated" / cfg["version"]
    out.mkdir(parents=True, exist_ok=True)
    summary = {}
    checksums = {}
    for split in SPLITS:
        targets = (
            cfg["category_targets"]
            if split == "train"
            else {c: cfg["holdout_decisions_per_category"] for c in cfg["category_targets"]}
        )
        cap = math.ceil(sum(targets.values()) * cfg["max_source_share"])
        selected = select(db, split, targets, cap)
        counts = {
            name: Counter() for name in ("category", "kind", "source", "label_origin", "license")
        }
        soft, total_tokens = 0, 0
        samples = defaultdict(list)
        path = out / (split + ".jsonl")
        with path.open("w") as stream:
            for fp in sorted(selected):
                record = json.loads(
                    db.execute(
                        "SELECT record FROM candidates WHERE fingerprint=?", (fp,)
                    ).fetchone()[0]
                )
                m = record["metadata"]
                m["split"] = split
                kind = record["input"]["questions"]["decision"]["type"]
                for name, value in (
                    ("category", m["category"]),
                    ("kind", kind),
                    ("source", m["upstream_source"]),
                    ("label_origin", m["label_origin"]),
                    ("license", m["license"]),
                ):
                    counts[name][value] += 1
                soft += int(sum(p > 1e-8 for p in record["target"]) > 1)
                total_tokens += m["payload_tokens"]
                stream.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n")
                if len(samples[m["category"]]) < 3:
                    samples[m["category"]].append(record)
        summary[split] = {
            "decisions": len(selected),
            **{k: dict(v) for k, v in counts.items()},
            "soft_target_decisions": soft,
            "mean_payload_tokens": total_tokens / max(1, len(selected)),
        }
        checksums[path.name] = sha256_file(path)
        (out / (split + "-review-samples.json")).write_text(
            json.dumps(samples, indent=2, ensure_ascii=False) + "\n"
        )
    stats["splits"] = summary
    stats["packed_requests"] = pack_release(root, cfg["version"], cfg, lock)
    for split in SPLITS:
        path = out / (split + "-requests.jsonl")
        checksums[path.name] = sha256_file(path)
        assert stats["packed_requests"][split]["decisions"] == summary[split]["decisions"]
    stats["training_shortfalls"] = {
        c: n - summary["train"]["category"].get(c, 0)
        for c, n in cfg["category_targets"].items()
        if n > summary["train"]["category"].get(c, 0)
    }
    stats["output_sha256"] = checksums
    stats["config_sha256"] = {
        str(p.relative_to(root)): sha256_file(p) for p in (root / "configs").glob("*.yaml")
    }
    stats["source_lock_sha256"] = sha256_file(root / "configs/sources.lock.json")
    stats["code_sha256"] = {
        str(p.relative_to(root)): sha256_file(p) for p in (root / "src/themis_data").glob("*.py")
    }
    stats["limitations"] = [
        "Category mappings are curated rules, not independent human annotations.",
        "Literal overlap checks do not detect all paraphrases or upstream pretraining exposure.",
        "Token budget is canonical request JSON plus 128 reserved tokens; final trainer serialization must be checked.",
        "Native held-out records are preserved; only the synthetic collection guarantees workflow-disjoint native splits.",
        "Source reuse classifications are metadata evidence, not independent legal clearance.",
    ]
    (out / "manifest.json").write_text(json.dumps(stats, indent=2, ensure_ascii=False) + "\n")
    (out / "sources.lock.json").write_text(json.dumps(lock, indent=2) + "\n")
    print(
        json.dumps(
            {
                "splits": {s: v["decisions"] for s, v in summary.items()},
                "shortfalls": stats["training_shortfalls"],
            },
            indent=2,
        ),
        flush=True,
    )
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=ROOT)
    args = parser.parse_args()
    root = args.root
    os.environ.setdefault("RAYON_NUM_THREADS", "4")
    cfg = yaml.safe_load((root / "configs/curation.yaml").read_text())
    mappings = yaml.safe_load((root / "configs/categories.yaml").read_text())
    lock = json.loads((root / "configs/sources.lock.json").read_text())
    for source in lock["sources"] + [lock["tokenizer"]]:
        for f in source["files"]:
            path = root / "data/raw" / source["repo"].replace("/", "--") / f["path"]
            if sha256_file(path) != f["sha256"]:
                raise ValueError("Checksum mismatch: " + str(path))
    stats = {"input_cases": 0, "benchmark_matches": []}
    db = store_candidates(root, lock, cfg, mappings, stats)
    export(root, db, cfg, lock, stats)
    db.close()


if __name__ == "__main__":
    main()