"""RAID loader (liamdugan/raid) — multi-domain, multi-generator AI-text data.

RAID is sorted by (domain, model, attack). Schema: text=`generation`, label=`model`
('human' = human, else = AI), `domain`, `attack` ('none' = clean).

IMPORTANT: HF only PARTIALLY converts RAID to parquet (it's huge), and because RAID is
domain-sorted the partial parquet contains only the *early* domains. So we never hardcode which
domains exist — `load_raid_records` returns domain-tagged records for whatever is present, and the
caller splits train/OOD from the domains actually observed. Shards are cached + column-pruned.
"""
from __future__ import annotations
import os
import random
import tempfile
from collections import Counter, defaultdict
import requests

_PARQUET_API = "https://datasets-server.huggingface.co/parquet?dataset=liamdugan/raid"
_CACHE_DIR = os.path.join(tempfile.gettempdir(), "raid_parquet")
_NEEDED = ["generation", "model", "domain", "attack"]


def _shard_files(split: str):
    meta = requests.get(_PARQUET_API, timeout=60).json()
    files = [f for f in meta.get("parquet_files", [])
             if f.get("split") == split and f.get("config") == "raid"]
    if not files:
        raise RuntimeError(f"No RAID parquet shards for split={split}; api keys={list(meta)}")
    files.sort(key=lambda f: f["url"])
    print(f"[raid] split='{split}': {len(files)} shard(s), partial-conversion="
          f"{requests.get(_PARQUET_API, timeout=60).json().get('partial')}")
    return files


def _download_shard(f) -> str:
    os.makedirs(_CACHE_DIR, exist_ok=True)
    path = os.path.join(_CACHE_DIR, f["split"] + "_" + f["filename"])
    if os.path.exists(path) and os.path.getsize(path) > 1000:
        return path
    print(f"[raid] downloading {f['filename']} ({f.get('size',0)/1e6:.0f} MB)...")
    b = requests.get(f["url"], timeout=900).content
    with open(path, "wb") as fh:
        fh.write(b)
    return path


def _read_shard(path):
    import pandas as pd
    try:
        import pyarrow.parquet as pq
        names = pq.read_schema(path).names
        cols = [c for c in _NEEDED if c in names]
        if len(cols) == len(_NEEDED):
            return pd.read_parquet(path, columns=cols)
    except Exception:
        pass
    return pd.read_parquet(path)


def load_raid_records(n_per_domain_class=300, split="train", attack="none",
                      min_words=30, max_words=600, max_shards=10, seed=42, verbose=True):
    """Return a list of {text, label(0=human,1=ai), domain, model} tagged records, balanced up to
    `n_per_domain_class` per (domain, label) across whatever domains the partial parquet contains."""
    counts = defaultdict(int)
    records = []
    raw_domains, raw_attacks = set(), set()
    for si, f in enumerate(_shard_files(split)[:max_shards]):
        df = _read_shard(_download_shard(f))
        for col in _NEEDED:
            if col not in df.columns:
                raise RuntimeError(f"RAID shard missing '{col}'; columns={list(df.columns)}")
        dom = df["domain"].astype(str).str.strip().str.lower()
        atk = df["attack"].astype(str).str.strip().str.lower()
        mdl = df["model"].astype(str).str.strip().str.lower()
        raw_domains.update(dom.unique().tolist()); raw_attacks.update(atk.unique().tolist())
        if si == 0 and verbose:
            print(f"[raid] shard0 shape={df.shape} domains={sorted(set(dom))[:12]} "
                  f"attacks={sorted(set(atk))[:6]}")
        mask = df["generation"].notna()
        if attack is not None:
            mask &= (atk == str(attack).lower())
        wc = df["generation"].str.split().str.len()
        mask &= (wc >= min_words) & (wc <= max_words)
        idx = df.index[mask].tolist()
        random.Random(seed + si).shuffle(idx)
        gen = df["generation"]
        for i in idx:
            d = dom[i]; label = 0 if mdl[i] == "human" else 1
            if counts[(d, label)] >= n_per_domain_class:
                continue
            counts[(d, label)] += 1
            records.append({"text": gen[i].strip(), "label": label, "domain": d, "model": mdl[i]})
        if verbose:
            print(f"[raid] after {f['filename']}: domains={sorted({k[0] for k in counts})}")
    if not records:
        raise RuntimeError(f"RAID collected 0 records. pre-filter domains={sorted(raw_domains)} "
                           f"attacks={sorted(raw_attacks)} attack-requested='{attack}'.")
    if verbose:
        print(f"[raid] records/domain: {dict(Counter(r['domain'] for r in records))}")
        print(f"[raid] label balance: {dict(Counter(r['label'] for r in records))}")
    return records


def load_raid(n_per_class=1500, split="train", domains=None, holdout_domains=None,
              attack="none", min_words=30, max_words=600, max_shards=10, seed=42, verbose=True):
    """(human, ai) convenience wrapper over load_raid_records, with optional domain include/holdout."""
    recs = load_raid_records(n_per_domain_class=max(50, n_per_class // 3), split=split, attack=attack,
                             min_words=min_words, max_words=max_words, max_shards=max_shards,
                             seed=seed, verbose=verbose)
    if domains:
        ds = {d.lower() for d in domains}; recs = [r for r in recs if r["domain"] in ds]
    if holdout_domains:
        hs = {d.lower() for d in holdout_domains}; recs = [r for r in recs if r["domain"] not in hs]
    human = [r["text"] for r in recs if r["label"] == 0]
    ai = [r["text"] for r in recs if r["label"] == 1]
    random.Random(seed + 2).shuffle(human); random.Random(seed + 3).shuffle(ai)
    return human[:n_per_class], ai[:n_per_class]