"""Build v6 SFT mix: FULL forensic set + chat/persona/tool/distill/SOP extras
+ raw TinyStories replay for fluency. Memory-frugal (streams the 2.1GB txt)."""
import json, random
from collections import Counter
from pathlib import Path
from data.tokenizer import load_tokenizer

rng = random.Random(20260802)
SEQ = 256
OUT = Path("data/sft_mix_v6.jsonl")

def load(p):
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]

def dedupe(rows):
    seen, out = set(), []
    for r in rows:
        k = (r.get("persona", "analyst"), r.get("user", "")[:180])
        if k in seen:
            continue
        seen.add(k); out.append(r)
    return out

def reservoir_sample(path, n, chunk=10000):
    rng2 = random.Random(7)
    keep = []
    with open(path, encoding="utf-8") as f:
        seen = 0
        for line in f:
            s = line.strip()
            if not s:
                continue
            seen += 1
            if len(keep) < n:
                keep.append(s)
            else:
                j = rng2.randrange(seen)
                if j < n:
                    keep[j] = s
    return keep

def visible(row, tok, u_id, a_id, eot):
    if "raw" in row:
        return True
    if not row.get("user") or not row.get("assistant"):
        return False
    p = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "none": ""}.get(row.get("persona", "analyst"), "<|analyst|>")
    p_ids = tok.encode(p).ids if row.get("persona", "analyst") != "none" else []
    ids = p_ids + [u_id] + tok.encode(row["user"]).ids + [a_id] + tok.encode(row["assistant"]).ids + [eot]
    return len(ids) <= SEQ

def main():
    tok = load_tokenizer("data/tokenizer.json")
    u_id = tok.token_to_id("<|user|>"); a_id = tok.token_to_id("<|assistant|>"); eot = tok.token_to_id("<|endoftext|>")
    v3 = dedupe(load("data/sft_mix_v3.jsonl"))
    truth = [r for r in v3 if r.get("user", "").startswith("Answer truthfully:")]
    chatish = [r for r in v3 if r.get("persona") == "analyst" and "raw" not in r and len(r.get("user", "")) < 90]

    mix = []
    mix += load("data/general_chat.jsonl")
    mix += load("data/persona_dialogue.jsonl")
    mix += load("data/tool_use.jsonl")
    mix += rng.sample(truth, 40)
    mix += rng.sample(chatish, 80)
    mix += rng.sample(load("data/sft_distill_mix.jsonl"), 160)
    mix += rng.sample(load("data/sft_sop_mix.jsonl"), 120)
    mix += dedupe(load("data/sft_forensic.jsonl"))  # ALL domain examples

    clean = [r for r in dedupe(mix) if visible(r, tok, u_id, a_id, eot)]
    for ln in reservoir_sample("data/TinyStoriesV2-GPT4-train.txt", 700):
        clean.append({"raw": ln, "persona": "none"})
    rng.shuffle(clean)
    with open(OUT, "w", encoding="utf-8") as f:
        for r in clean:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("total", len(clean), dict(Counter(r.get("persona", "?") for r in clean)), flush=True)

if __name__ == "__main__":
    main()