#!/usr/bin/env python3
"""Build the static shards consumed by the WildChat-AQA eval-tab viewer.

WildChat-AQA (https://huggingface.co/datasets/wtzhang-nlp/wildchat_aqa) is an
*aggregative* multiple-choice QA benchmark over WildChat chat logs: each
question asks for an aggregate statistic (e.g. "most popular topics for user X")
answered by aggregating over the conversations matching a *condition*.

The QA parquet has no supporting documents, so this script reconstructs them:
it links every question to the WildChat conversations satisfying its
``condition_type``/``condition_value`` (from ``wildchat_aqa_conversations``),
resolves label ids via the repo's ``wildchat_aqa_taxonomy/``, joins the PROBE
``generated_query`` (from the ``…with_embedding_and_gpt_generated_query``
dataset), and emits:

* ``corpus/<conv_hash>.json``   one shard per referenced conversation
  (summary + metadata + full raw turns), lazy-loaded by the UI.
* ``eval/records/<qhash>.json`` per question: question, resolved condition,
  weighted options (answer = argmax weight), resolved ``target_candidates``
  aggregation evidence, ``generated_query``, and capped supporting-conv refs.
* ``eval/index.json``           light list for the sidebar (search + filters).

Condition matching: **OR within a repeated condition type, AND across types**
(so ``(user_name, user_name)`` = either user; ``(country, label_level_1)`` =
country AND topic). ``time_week`` matches a 7-day window; label values encode
as ``"<L1>.<L2>"``.

Run from the viewer repo root::

    python scripts/build_wildchat_aqa.py --src /tmp/wc_src --out . --cap 30
"""
from __future__ import annotations

import argparse
import collections
import datetime
import glob
import json
import os
import sys
from pathlib import Path

import pyarrow.parquet as pq

# raw conversation-turn fields worth keeping in the corpus shard
TURN_FIELDS = ("role", "content", "country", "language", "timestamp")


def week_start(ts) -> str | None:
    if ts is None:
        return None
    if hasattr(ts, "date"):
        d = ts.date()
    else:
        d = datetime.datetime.fromisoformat(str(ts)).date()
    return (d - datetime.timedelta(days=d.weekday())).isoformat()


def load_taxonomy(src: Path):
    l1 = {c["index"]: c["class_name"]
          for c in json.loads((src / "taxonomy" / "step_3_3_manual_level_1_taxonomy_result.json").read_text())["classes"]}
    l2 = {}  # "L1.L2" -> name
    for f in glob.glob(str(src / "taxonomy" / "step_4_4_manual_level_2_taxonomy_result_*.json")):
        idx = int(os.path.basename(f).split("_")[-1].split(".")[0])
        for j, c in enumerate(json.loads(Path(f).read_text())["classes"]):
            l2[f"{idx}.{j}"] = c["class_name"]
    return l1, l2


def load_generated_query(src: Path) -> dict:
    p = src / "generated_query.jsonl"
    out = {}
    if not p.exists():
        return out
    for line in p.open():
        d = json.loads(line)
        raw = d.get("generated_query")
        try:
            out[d["hash"]] = json.loads(raw) if isinstance(raw, str) else raw
        except (json.JSONDecodeError, TypeError):
            out[d["hash"]] = {"explanation": raw} if raw else None
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--src", default="/tmp/wc_src", help="dir holding qa.parquet, conversations/, taxonomy/, generated_query.jsonl")
    ap.add_argument("--out", default=".", help="viewer repo root to write shards into")
    ap.add_argument("--cap", type=int, default=30, help="max supporting conversations per question")
    ap.add_argument("--tc-cap", type=int, default=50, help="max target_candidates shown per question")
    args = ap.parse_args()

    src = Path(args.src)
    out = Path(args.out)
    if not (src / "qa.parquet").exists():
        print(f"ERROR: {src/'qa.parquet'} not found", file=sys.stderr)
        return 1

    l1, l2 = load_taxonomy(src)
    gq = load_generated_query(src)
    print(f"taxonomy: {len(l1)} L1 / {len(l2)} L2 names | generated_query: {len(gq)}")

    def resolve_label(target_type: str, value: str) -> str | None:
        if target_type == "label_level_1":
            try:
                return l1.get(int(value))
            except (ValueError, TypeError):
                return None
        if target_type == "label_level_2":
            return l2.get(str(value))
        return None

    # ── pass over conversations: meta + trimmed raw turns + inverted indices ──
    meta: dict = {}
    raw_turns: dict = {}
    idx = {k: collections.defaultdict(set) for k in
           ("user_name", "label_level_1", "label_level_2", "country", "language", "keywords_aggregated", "time_week")}
    cols = ["hash", "user_name", "classes_level_1", "classes_level_2",
            "keywords_aggregated", "keywords", "timestamp", "token_count", "summary", "conversation"]
    for f in sorted(glob.glob(str(src / "conversations" / "*.parquet"))):
        for r in pq.read_table(f, columns=cols).to_pylist():
            h = r["hash"]
            conv = r["conversation"] or []
            c0 = conv[0] if conv else {}
            l2set = set()
            for i, subs in enumerate(r["classes_level_2"] or []):
                for s in (subs or []):
                    l2set.add(f"{i}.{s}")
            l1set = set(r["classes_level_1"] or [])
            kwv = set(k["value"] for k in (r["keywords_aggregated"] or []) if k.get("value"))
            wk = week_start(r["timestamp"])
            meta[h] = {
                "user_name": r["user_name"], "l1": l1set, "l2": l2set,
                "country": c0.get("country"), "language": c0.get("language"),
                "week": wk, "tok": r["token_count"], "summary": r["summary"],
                "timestamp": str(r["timestamp"]) if r["timestamp"] else None,
                "keywords": r["keywords"] or [], "kw_agg": r["keywords_aggregated"] or [],
            }
            raw_turns[h] = [{k: t.get(k) for k in TURN_FIELDS} for t in conv]
            # indices
            if r["user_name"]:
                idx["user_name"][r["user_name"]].add(h)
            for x in l1set:
                idx["label_level_1"][str(x)].add(h)
            for x in l2set:
                idx["label_level_2"][x].add(h)
            if c0.get("country"):
                idx["country"][c0["country"]].add(h)
            if c0.get("language"):
                idx["language"][c0["language"]].add(h)
            for x in kwv:
                idx["keywords_aggregated"][x].add(h)
            if wk:
                idx["time_week"][wk].add(h)
    print(f"conversations: {len(meta)}")

    all_hashes = set(meta.keys())

    def support(cond_types, cond_vals) -> set:
        groups = collections.defaultdict(list)
        for t, v in zip(cond_types, cond_vals):
            groups[t].append(v)
        result = None
        for t, vals in groups.items():
            u = set()
            for v in vals:
                if t == "time_week":
                    u |= idx["time_week"].get(week_start(datetime.datetime.fromisoformat(str(v))), set())
                else:
                    u |= idx[t].get(str(v), set())
            result = u if result is None else (result & u)
        return result if result is not None else all_hashes

    # ── QA: build records + collect referenced conversations ──
    rec_dir = out / "eval" / "records"
    corpus_dir = out / "corpus"
    rec_dir.mkdir(parents=True, exist_ok=True)
    corpus_dir.mkdir(parents=True, exist_ok=True)

    rows = pq.read_table(str(src / "qa.parquet")).to_pylist()
    index = []
    referenced: set = set()

    for r in rows:
        qh = r["hash"]
        ctypes, cvals = list(r["condition_type"]), list(r["condition_value"])
        s = support(ctypes, cvals)
        total = len(s)
        across_all = len(ctypes) == 0
        chosen = sorted(s, key=lambda h: (-(meta[h]["tok"] or 0), h))[: args.cap]
        referenced.update(chosen)

        # answer = argmax weight
        weights = list(r["option_weights"])
        opts = list(r["options"])
        best_i = max(range(len(weights)), key=lambda i: weights[i]) if weights else -1

        condition = [{"type": t, "value": v, "value_name": resolve_label(t, v)}
                     for t, v in zip(ctypes, cvals)]
        tcs = r["target_candidates"][: args.tc_cap]
        target_candidates = [{"value": c["value"], "count": c["count"],
                              "value_name": resolve_label(r["target_type"], c["value"])} for c in tcs]

        support_refs = [{
            "hash": h, "user_name": meta[h]["user_name"], "summary": meta[h]["summary"],
            "country": meta[h]["country"], "language": meta[h]["language"],
            "week": meta[h]["week"], "token_count": meta[h]["tok"],
        } for h in chosen]

        record = {
            "hash": qh,
            "question": r["question"],
            "target_type": r["target_type"],
            "keywords_type": r["keywords_type"],
            "condition": condition,
            "options": [{"text": opts[i], "weight": weights[i]} for i in range(len(opts))],
            "answer": {"text": opts[best_i], "weight": weights[best_i]} if best_i >= 0 else None,
            "target_candidates": target_candidates,
            "target_candidates_total": len(r["target_candidates"]),
            "generated_query": gq.get(qh),
            "n_support_total": total,
            "n_support_shown": len(chosen),
            "capped": total > len(chosen),
            "across_all": across_all,
            "support": support_refs,
        }
        (rec_dir / f"{qh}.json").write_text(json.dumps(record, ensure_ascii=False, default=str))

        index.append({
            "hash": qh,
            "question": r["question"][:200],
            "target_type": r["target_type"],
            "keywords_type": r["keywords_type"],
            "n_conditions": len(condition),
            "cond_types": sorted({c["type"] for c in condition}),
            "condition_summary": ", ".join(f"{c['type']}={c['value_name'] or c['value']}" for c in condition) or "(across all)",
            "n_support": total,
        })

    (out / "eval" / "index.json").write_text(json.dumps(index, ensure_ascii=False, default=str))

    # ── corpus shards for referenced conversations ──
    for h in referenced:
        m = meta[h]
        l1_names = sorted({l1.get(x) for x in m["l1"] if l1.get(x)})
        l2_names = sorted({l2.get(x) for x in m["l2"] if l2.get(x)})
        (corpus_dir / f"{h}.json").write_text(json.dumps({
            "hash": h,
            "user_name": m["user_name"],
            "timestamp": m["timestamp"],
            "country": m["country"],
            "language": m["language"],
            "token_count": m["tok"],
            "summary": m["summary"],
            "keywords": m["keywords"],
            "keywords_aggregated": m["kw_agg"],
            "classes_level_1": l1_names,
            "classes_level_2": l2_names,
            "n_turns": len(raw_turns[h]),
            "conversation": raw_turns[h],
        }, ensure_ascii=False, default=str))

    print(f"records: {len(rows)} | referenced convs: {len(referenced)} | index entries: {len(index)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())