"""
Unify any raw file to standard pretrain schema:
  text, lang, source, domain, safety_tag

Handles known mismatched schemas found on ViuAI/viu-mini-raw-pretrain:
  - wikipedia (title/text/url)            -> text = title + text
  - samanantar (src/tgt/translation)      -> text = "English: ...\\nHindi: ..."
  - gsm8k/math (question/answer)          -> text = Q + A with <soch> wrapper
  - dolly/alpaca/sharegpt/lima (instruction/response ...) -> text templated
  - grammar (text/topic/difficulty)       -> domain = hindi_vyakaran_<topic>
  - sciq/textbooks (question/choices/answer, text/context)
  - toxicity/civil_comments (text/toxicity labels) -> safety_tag=toxic
  - uncensored/*                          -> safety_tag=uncensored (rakhna hai, delete nahi)
  - csv/jsonl/txt/parquet auto-detect

Usage:
  python data/scripts/standardize_to_unified.py --in science/sciq_train.parquet --source sciq --lang english --domain science_qa --safety safe --out /tmp/out.parquet
  python data/scripts/standardize_to_unified.py --in hinglish/dolly15k-hinglish.parquet --auto --out /tmp/out.parquet --dry-run

Kaggle/HF repair me isko streaming + sharded mode me chalao (see HF_REPAIR_RUNBOOK).
"""
import argparse
import hashlib
import json
import sys
from pathlib import Path

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

UNIFIED_COLS = ["text", "lang", "source", "domain", "safety_tag"]

DEFAULT_SAFETY_BY_FOLDER = {
    "toxicity": "toxic",
    "uncensored": "uncensored",
}


def infer_safety(source: str, domain: str, folder: str = "") -> str:
    s = (source + " " + domain + " " + folder).lower()
    if "uncensor" in s or "unfilter" in s or "vortex" in s or "alpaca-uncens" in s:
        return "uncensored"
    if "toxic" in s or "hate" in s or "offensive" in s or "profan" in s or "civil_comments" in s or "hasoc" in s or "hatecheck" in s or "prism" in s:
        return "toxic"
    return "safe"


def row_to_text(rec: dict) -> str | None:
    """Map any known record shape to a single text string. Returns None if unusable."""
    if not isinstance(rec, dict):
        return str(rec).strip() or None
    # already unified
    if rec.get("text") and isinstance(rec["text"], str) and len(rec["text"].strip()) >= 5:
        return rec["text"].strip()
    # wikipedia: title + text
    if "title" in rec and "text" in rec:
        t = f"{rec.get('title','').strip()}\n\n{(rec.get('text') or '').strip()}".strip()
        return t or None
    # samanantar / translation variants
    if "translation" in rec and isinstance(rec["translation"], dict):
        tr = rec["translation"]
        en = (tr.get("en") or tr.get("src") or "").strip()
        hi = (tr.get("hi") or tr.get("tgt") or "").strip()
        if en and hi:
            return f"English: {en}\nHindi: {hi}"
    for k in ("src", "source_text", "source_lang"):
        pass
    if "en" in rec and "hi" in rec and isinstance(rec["en"], str):
        en, hi = rec["en"].strip(), rec["hi"].strip()
        if en and hi:
            return f"English: {en}\nHindi: {hi}"
    # math QA
    if "question" in rec and "answer" in rec:
        q, a = str(rec["question"]).strip(), str(rec["answer"]).strip()
        if q and a:
            return f"<|user|> {q} <|assistant|> <soch> {a} </soch> {a}"
    if "problem" in rec and "solution" in rec:
        q, a = str(rec["problem"]).strip(), str(rec["solution"]).strip()
        if q and a:
            return f"<|user|> {q} <|assistant|> <soch> {a} </soch> {a}"
    # instruct variants: instruction/response, prompt/completion, input/output
    instr = rec.get("instruction") or rec.get("prompt") or rec.get("input") or rec.get("instruction_en") or rec.get("instruction_hien")
    resp = rec.get("response") or rec.get("output") or rec.get("completion") or rec.get("response_en") or rec.get("response_hien")
    if instr and resp:
        # hinglish dolly has both en + hien -> prefer hien, keep en as prefix if different
        hien_i = rec.get("instruction_hien")
        hien_r = rec.get("response_hien")
        if hien_i and hien_r:
            return f"<|user|> {str(hien_i).strip()} <|assistant|> {str(hien_r).strip()}"
        return f"<|user|> {str(instr).strip()} <|assistant|> {str(resp).strip()}"
    # sciq: question + distractors + correct
    if "question" in rec and ("correct_answer" in rec or "support" in rec):
        parts = [str(rec.get("question", "")).strip()]
        for k in ("distractor1", "distractor2", "distractor3", "support"):
            if rec.get(k):
                parts.append(str(rec[k]).strip())
        if "correct_answer" in rec and rec["correct_answer"]:
            parts.append(f"Answer: {str(rec['correct_answer']).strip()}")
        t = "\n".join([p for p in parts if p])
        return t or None
    # generic: concatenate all string values
    vals = [str(v).strip() for v in rec.values() if isinstance(v, str) and len(str(v).strip()) >= 10]
    if vals:
        # avoid exploding wide tables (civil_comments has many float label cols - already filtered since not str)
        t = "\n".join(vals[:4])
        return t[:20000] or None
    return None


def standardize_record(rec: dict, fallback_lang: str, fallback_source: str, fallback_domain: str, fallback_safety: str) -> dict | None:
    text = row_to_text(rec)
    if not text or len(text.strip()) < 10 or len(text.split()) < 3:
        return None
    lang = rec.get("lang") if isinstance(rec.get("lang"), str) and rec.get("lang") else fallback_lang
    source = rec.get("source") if isinstance(rec.get("source"), str) and rec.get("source") else fallback_source
    domain = rec.get("domain") if isinstance(rec.get("domain"), str) and rec.get("domain") else fallback_domain
    # grammar files carry topic -> fold into domain
    if rec.get("topic") and "vyakaran" not in domain and "grammar" not in domain:
        domain = f"{domain}_{rec.get('topic')}" if domain else str(rec.get("topic"))
    safety = rec.get("safety_tag") if isinstance(rec.get("safety_tag"), str) and rec.get("safety_tag") else fallback_safety
    if safety == "safe":
        # auto-infer toxic/uncensored so they are TAGGED not lost (user wants to keep them)
        safety = infer_safety(str(source), str(domain))
    return {"text": text.strip(), "lang": str(lang), "source": str(source), "domain": str(domain), "safety_tag": str(safety)}


def iter_input(path: Path):
    suf = path.suffix.lower()
    if suf == ".parquet":
        import pyarrow.parquet as pq
        pf = pq.ParquetFile(str(path))
        for batch in pf.iter_batches(batch_size=10000):
            tbl = batch.to_pylist()
            for r in tbl:
                yield r
    elif suf == ".jsonl" or (suf == ".json" and "jsonl" in path.name):
        with open(path, encoding="utf-8", errors="ignore") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    yield json.loads(line)
                except Exception:
                    yield {"text": line}
    elif suf == ".json":
        with open(path, encoding="utf-8", errors="ignore") as f:
            data = json.load(f)
        if isinstance(data, list):
            for r in data:
                yield r if isinstance(r, dict) else {"text": str(r)}
        elif isinstance(data, dict):
            yield data
    elif suf == ".csv":
        import csv
        with open(path, encoding="utf-8", errors="ignore", newline="") as f:
            for r in csv.DictReader(f):
                yield r
    elif suf == ".txt":
        with open(path, encoding="utf-8", errors="ignore") as f:
            for line in f:
                if line.strip():
                    yield {"text": line.strip()}
    else:
        raise ValueError(f"unsupported input: {suf}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--out", required=True, help="output parquet (unified schema)")
    ap.add_argument("--source", default="unknown")
    ap.add_argument("--lang", default="hinglish")
    ap.add_argument("--domain", default="general")
    ap.add_argument("--safety", default="safe")
    ap.add_argument("--auto", action="store_true", help="infer source/lang/domain/safety from input path")
    ap.add_argument("--dry-run", action="store_true", help="only print first 3 converted rows, no write")
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    inp = Path(args.inp)
    src, lang, dom, safety = args.source, args.lang, args.domain, args.safety
    if args.auto:
        parts = inp.as_posix().split("/")
        folder = parts[-2] if len(parts) >= 2 else ""
        stem = inp.stem
        src = folder + "_" + stem if folder else stem
        # lang guess
        low = (folder + " " + stem).lower()
        if "vyakaran" in low or "sandhi" in low or "samas" in low:
            lang = "hindi"
        elif "hindi" in low and "hinglish" not in low:
            lang = "hindi"
        elif "hinglish" in low or "roman" in low:
            lang = "hinglish"
        elif "english" in low or "wiki" in low or "sciq" in low or "textbook" in low:
            lang = "english"
        dom = folder or "general"
        safety = infer_safety(src, dom, folder)

    kept, dropped = 0, 0
    sample = []
    seen = set()
    out_rows = []
    for i, rec in enumerate(iter_input(inp)):
        if args.limit and kept >= args.limit:
            break
        std = standardize_record(rec, lang, src, dom, safety)
        if not std:
            dropped += 1
            continue
        h = hashlib.blake2b(std["text"].encode("utf-8"), digest_size=8).digest()
        if h in seen:
            dropped += 1
            continue
        seen.add(h)
        kept += 1
        if args.dry_run and len(sample) < 3:
            sample.append(std)
        elif not args.dry_run:
            out_rows.append(std)

    print(f"[unify] {inp} -> kept={kept:,} dropped={dropped:,} lang={lang} source={src} domain={dom} safety={safety}")
    if args.dry_run:
        for r in sample:
            print("---", {k: (str(v)[:160]) for k, v in r.items()})
        return
    import pandas as pd
    df = pd.DataFrame(out_rows, columns=UNIFIED_COLS)
    outp = Path(args.out)
    outp.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(outp, index=False, compression="snappy")
    print(f"[ok] wrote {outp} ({outp.stat().st_size/1024/1024:.1f} MB, {len(df):,} rows)")

if __name__ == "__main__":
    main()