"""Loaders for LocalLLaMA/typed-decisions and Praveenrajus/jev-bench, plus deterministic manifests."""
from __future__ import annotations

import hashlib
import json
import math
import random
import re
from collections import defaultdict
from pathlib import Path

from .schema import Example, Question, align_distribution

TYPED_DECISIONS = "LocalLLaMA/typed-decisions"
JEV_BENCH = "Praveenrajus/jev-bench"
JEV_BENCH_CONFIGS = [
    "arc_challenge", "boolq", "chaosnli", "civil_comments", "fever_evidence", "go_emotions",
    "helpsteer2_helpfulness", "helpsteer2_verbosity", "massive", "measuring_hate_speech", "mmlu",
    "mnli", "paws", "sms_spam", "sst5", "strategyqa_closed", "strategyqa_grounded", "stsb", "yelp5",
]
TRUE_LABELS = {"1", "true", "yes"}
FALSE_LABELS = {"0", "false", "no"}


def _questions_from_map(qmap: dict) -> list[Question]:
    return [Question.from_jev(qid, spec) for qid, spec in qmap.items()]


def _label_option(label, q: Question) -> str | None:
    """The row's discrete gold label as one of q's option names, or None when it is missing or not an option.

    A boolean names the noul option "true" or "false", and a number its string form (a score level such as "2").
    """
    if isinstance(label, bool):
        label = "true" if label else "false"
    if label is None or str(label) not in q.options:
        return None
    return str(label)


def typed_decisions_example(row: dict) -> Example:
    questions = _questions_from_map(json.loads(row["questions"]))
    gold_raw = json.loads(row["gold"])
    gold, labels = {}, {}
    for q in questions:
        entry = gold_raw[q.qid]
        gold[q.qid] = align_distribution(entry["probabilities"], q.options)
        # The card's discrete label breaks ties in the distribution; hev.evaluate scores accuracy against it.
        label = _label_option(entry.get("label"), q)
        if label is not None:
            labels[q.qid] = label
    meta = {"workflow": row["workflow"]}
    if labels:
        meta["labels"] = labels
    return Example(id=row["id"], source=f"typed-decisions/{row['workflow']}",
                   state=row["state"], questions=questions, gold=gold, meta=meta)


def _positional_distribution(soft, k: int) -> list[float] | None:
    if len(soft) != k:
        return None
    if not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in soft):
        return None
    vals = [float(v) for v in soft]
    if any(v < 0 or not math.isfinite(v) for v in vals):
        return None
    total = sum(vals)
    if total <= 0:
        return None
    return [v / total for v in vals]


def jev_bench_example(row: dict, config: str) -> Example:
    q = Question.from_jev("q", json.loads(row["question"]))
    soft = row.get("soft_label")
    if isinstance(soft, str):
        soft = json.loads(soft)
    dist = None
    if isinstance(soft, dict) and set(soft) & set(q.options):
        dist = align_distribution(soft, q.options)
    elif (isinstance(soft, (int, float)) and not isinstance(soft, bool) and math.isfinite(soft)
          and q.type == "noul"):
        p = min(1.0, max(0.0, float(soft)))
        dist = [p if o == "true" else 1.0 - p for o in q.options]
    elif isinstance(soft, (list, tuple)) and q.type == "score":
        dist = _positional_distribution(soft, q.k)
    if dist is None:
        label = str(row["label"])
        if q.type == "noul":
            norm = label.strip().lower()
            if norm in TRUE_LABELS:
                label = "true"
            elif norm in FALSE_LABELS:
                label = "false"
            else:
                raise ValueError(f"{row['id']}: unrecognised noul label {row['label']!r}")
        if label not in q.options:
            raise ValueError(f"{row['id']}: label {label!r} not in options {q.options}")
        dist = [1.0 if o == label else 0.0 for o in q.options]
    meta = row.get("meta") or "{}"
    meta = json.loads(meta) if isinstance(meta, str) else dict(meta)
    meta["config"] = config
    return Example(id=row["id"], source=f"jev-bench/{config}", state=row["state"],
                   questions=[q], gold={"q": dist}, meta=meta)


def split_validation(examples: list[Example], frac: float = 0.1, seed: int = 42):
    by_source = defaultdict(list)
    for ex in examples:
        by_source[ex.source].append(ex)
    train, val = [], []
    for source in sorted(by_source):
        group = sorted(by_source[source], key=lambda e: e.id)
        rng = random.Random(f"{seed}:{source}")
        rng.shuffle(group)
        n_val = max(1, round(len(group) * frac))
        val.extend(group[:n_val])
        train.extend(group[n_val:])
    return train, val


def write_manifest(examples: list[Example], path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for ex in examples:
            f.write(json.dumps(ex.to_dict(), ensure_ascii=False) + "\n")


def read_manifest(path) -> list[Example]:
    with Path(path).open("r", encoding="utf-8") as f:
        return [Example.from_dict(json.loads(line)) for line in f if line.strip()]


def manifest_checksum(path) -> str:
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()[:16]


def _is_missing_split(err: Exception, split: str) -> bool:
    """True for the error `datasets` raises when a config has no such split (the only load error that is skipped)."""
    msg = str(err)
    return ("Unknown split" in msg or "Bad split" in msg) and split in msg


def _dataset_revision(ds) -> str | None:
    """The Hub commit a loaded split came from, read off its cache path (.../<config>/<version>/<sha>/...)."""
    for f in getattr(ds, "cache_files", None) or []:
        name = Path(str(f.get("filename", ""))).parent.name
        if re.fullmatch(r"[0-9a-f]{40}", name):
            return name
    return None


def _sample(rows, k: int, seed: int, key: str):
    idx = list(range(len(rows)))
    random.Random(f"{seed}:{key}").shuffle(idx)
    return [rows[i] for i in sorted(idx[:k])]


def build_manifests(out_dir, seed: int = 42, per_config: dict | None = None) -> dict:
    # Imported here, inside try/except, so neither the tests nor Hub users of the flattened data.py need `datasets`:
    # transformers' check_imports skips imports inside try blocks.
    try:
        from datasets import load_dataset
    except ImportError as err:
        raise ImportError('building the manifests needs the `datasets` package: pip install datasets '
                          '(or pip install -e ".[data]")') from err

    per_config = per_config or {"train": 300, "validation": 100, "test": 200}
    out = Path(out_dir)
    revisions = defaultdict(set)

    def load(repo, config, split):
        ds = load_dataset(repo, config, split=split)
        sha = _dataset_revision(ds)
        if sha:
            revisions[repo].add(sha)
        return ds

    td_train = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "train")]
    td_test = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "test")]
    train, val = split_validation(td_train, 0.1, seed)
    jb_test = []
    for config in JEV_BENCH_CONFIGS:
        for split, target in (("train", train), ("validation", val), ("test", jb_test)):
            try:
                rows = load(JEV_BENCH, config, split)
            except ValueError as err:
                if not _is_missing_split(err, split):
                    raise
                print(f"skip {config}/{split}: {err}", flush=True)
                continue
            rows = [rows[i] for i in range(len(rows))]
            for r in _sample(rows, per_config[split], seed, f"{config}/{split}"):
                target.append(jev_bench_example(r, config))
    write_manifest(train, out / "train.jsonl")
    write_manifest(val, out / "validation.jsonl")
    write_manifest(td_test, out / "test_td.jsonl")
    write_manifest(jb_test, out / "test_jb.jsonl")
    counts = {"train": len(train), "validation": len(val), "test_td": len(td_test), "test_jb": len(jb_test)}
    for name in ("train", "validation", "test_td", "test_jb"):
        counts[f"{name}_sha"] = manifest_checksum(out / f"{name}.jsonl")
    if revisions:  # the dataset commit each repo was loaded from (a sorted list if the cache mixed several)
        counts["revisions"] = {repo: next(iter(shas)) if len(shas) == 1 else sorted(shas)
                               for repo, shas in sorted(revisions.items())}
    (out / "counts.json").write_text(json.dumps(counts, indent=2))
    return counts


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()
    print(json.dumps(build_manifests(args.out, args.seed), indent=2))