"""
export_distillation.py — turn captured teacher examples into student-ready JSONL.

Offline/admin tool. Reads the `distillation_examples` Firestore collection (written by
distill.py) and emits one JSONL file per task under an output directory:

  intent_parse.jsonl / translate.jsonl / codegen.jsonl / coach.jsonl / insight.jsonl
      → chat-format SFT  {"messages":[{system},{user},{assistant}]}
  vision.jsonl          → {"image_uri", "instruction", "output"}
  asr.jsonl             → {"audio_uri", "transcript"}
  rejected.jsonl        → hard negatives (verdict=rejected), for DPO / filtering

This is sequence-level (response-based) distillation: targets are the teacher's text /
JSON outputs and the human-verified gold label — Gemini exposes no logits.

Usage:
  python scripts/export_distillation.py --out ./dataset \
      --verdicts confirmed corrected pending
  (defaults: SFT uses confirmed+corrected; pass --include-pending to add unlabeled)

Env: FIREBASE (service-account JSON), as used by the app.
"""

import argparse
import json
import os
from collections import Counter, defaultdict
from typing import Any, Dict, List, Optional


def _init_db():
    import firebase_admin
    from firebase_admin import credentials, firestore
    if not firebase_admin._apps:
        sa = json.loads(os.environ["FIREBASE"])
        cred = credentials.Certificate(sa)
        opts = {}
        bucket = os.environ.get("FIREBASE_STORAGE_BUCKET")
        if bucket:
            opts["storageBucket"] = bucket
        firebase_admin.initialize_app(cred, opts)
    return firestore.client()


# Tasks that train as text chat SFT.
_TEXT_TASKS = {"intent_parse", "translate", "codegen", "coach", "insight"}


def _gold_or_teacher(ex: Dict[str, Any]) -> Optional[str]:
    """Prefer the human-corrected gold output; fall back to the teacher output."""
    return ex.get("gold_output") or ex.get("output_raw")


def _to_chat(ex: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    inp = ex.get("input") or {}
    user = inp.get("text") or inp.get("transcript") or ""
    system = inp.get("system_prompt") or ""
    target = _gold_or_teacher(ex)
    if not user or not target:
        return None
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": user})
    messages.append({"role": "assistant", "content": target})
    return {"messages": messages, "task": ex.get("task"),
            "teacher_model": ex.get("teacher_model"),
            "prompt_version": ex.get("prompt_version"),
            "verdict": ex.get("verdict")}


def _to_vision(ex: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    inp = ex.get("input") or {}
    uri = inp.get("media_uri")
    target = _gold_or_teacher(ex)
    if not uri or not target:
        return None
    return {"image_uri": uri, "instruction": inp.get("caption") or "Extract the bookkeeping transactions as JSON.",
            "output": target, "verdict": ex.get("verdict"),
            "teacher_model": ex.get("teacher_model")}


def _to_asr(ex: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    inp = ex.get("input") or {}
    uri = inp.get("media_uri")
    transcript = ex.get("output_raw")
    if not uri or not transcript:
        return None
    return {"audio_uri": uri, "transcript": transcript, "teacher_model": ex.get("teacher_model")}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="./dataset", help="output directory")
    ap.add_argument("--verdicts", nargs="*", default=["confirmed", "corrected"],
                    help="verdicts to include for SFT (default: confirmed corrected)")
    ap.add_argument("--include-pending", action="store_true",
                    help="also include unlabeled (pending) examples in SFT")
    ap.add_argument("--limit", type=int, default=0, help="max docs (0 = all)")
    args = ap.parse_args()

    sft_verdicts = set(args.verdicts)
    if args.include_pending:
        sft_verdicts.add("pending")

    db = _init_db()
    os.makedirs(args.out, exist_ok=True)

    writers: Dict[str, Any] = {}

    def _w(name: str):
        if name not in writers:
            writers[name] = open(os.path.join(args.out, f"{name}.jsonl"), "w", encoding="utf-8")
        return writers[name]

    counts: Counter = Counter()
    # Dedup identical (user,target) pairs per task.
    seen: Dict[str, set] = defaultdict(set)

    q = db.collection("distillation_examples")
    docs = q.limit(args.limit).stream() if args.limit else q.stream()

    for doc in docs:
        ex = doc.to_dict() or {}
        task = ex.get("task") or "other"
        verdict = ex.get("verdict") or "pending"

        if verdict == "rejected":
            rec = _to_chat(ex) or {"input": ex.get("input"), "output": ex.get("output_raw")}
            _w("rejected").write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
            counts["rejected"] += 1
            continue

        if task == "asr":
            rec = _to_asr(ex)
            if rec:
                _w("asr").write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
                counts["asr"] += 1
            continue

        if verdict not in sft_verdicts:
            continue

        if task == "vision":
            rec = _to_vision(ex)
            if rec:
                _w("vision").write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
                counts["vision"] += 1
            continue

        if task in _TEXT_TASKS:
            rec = _to_chat(ex)
            if not rec:
                continue
            key = (rec["messages"][-2]["content"], rec["messages"][-1]["content"])
            if key in seen[task]:
                continue
            seen[task].add(key)
            _w(task).write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
            counts[task] += 1

    for f in writers.values():
        f.close()

    print("Exported examples per file:")
    for name, n in sorted(counts.items()):
        print(f"  {name}.jsonl: {n}")
    print(f"Output: {os.path.abspath(args.out)}")


if __name__ == "__main__":
    main()