#!/opt/homebrew/bin/python3
"""Export gabgod.line_bank -> SFT JSONL (mlx_lm.lora chat format) for retraining.

Only exports approved lines by default (--all to include unreviewed).
Writes sft/train.jsonl + sft/valid.jsonl (95/5 deterministic split).

  ./export_sft.py            # approved only
  ./export_sft.py --all      # everything (pre-review experiments)
"""
import argparse, json, hashlib
from pathlib import Path
import psycopg2

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "config.json").read_text())
PERSONAS = json.loads((HERE / "personas.json").read_text())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--all", action="store_true")
    args = ap.parse_args()

    conn = psycopg2.connect(CFG["db"])
    cur = conn.cursor()
    where = "" if args.all else "where approved is true"
    cur.execute(f"select persona, situation, line from gabgod.line_bank {where}")

    out = HERE / "sft"
    out.mkdir(exist_ok=True)
    train = open(out / "train.jsonl", "w")
    valid = open(out / "valid.jsonl", "w")
    n = {"train": 0, "valid": 0}
    for persona, situation, line in cur.fetchall():
        rec = {"messages": [
            {"role": "system",
             "content": f"You are {persona.replace('_',' ')} in a 90s Australian shopping town. "
                        + PERSONAS.get(persona, {}).get("voice", "")},
            {"role": "user", "content": situation},
            {"role": "assistant", "content": line}]}
        bucket = "valid" if int(hashlib.sha1(line.encode()).hexdigest(), 16) % 20 == 0 else "train"
        (valid if bucket == "valid" else train).write(json.dumps(rec) + "\n")
        n[bucket] += 1
    print(f"wrote {n['train']} train / {n['valid']} valid -> {out}/")
    if not args.all and sum(n.values()) == 0:
        print("0 approved lines — review first (./gabgod review) or use --all")


if __name__ == "__main__":
    main()