"""Build v8 SFT mix: ALL forensic + doubled claim-specific examples + chat
extras + raw replay. Goal: strengthen input->verdict conditioning so the
template doesn't dominate."""
import json, random, re
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_v8.jsonl")
STOP = {"identify","factual","assertion","compare","against","known","records","note","missing","context",
        "statement","claim","source","beyond","itself","verdict","confidence","reasoning","checklist","evidence",
        "requires","about","with","that","this","what","their","they","will","have","from","your","which","there","its",
        "who","how","and","the","are","for","not","but","was","were","been","has","had","one","two","any","can","could"}

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

def words(s):
    return set(re.findall(r"[A-Za-z]{5,}", s.lower()))

def claim_specific(r):
    u = words(r.get("user", "")); a = words(r.get("assistant", ""))
    return bool((u & a) - STOP)

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 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 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|>")
    forensic = dedupe(load("data/sft_forensic.jsonl"))
    spec = [r for r in forensic if claim_specific(r)]
    print(f"forensic {len(forensic)} claim-specific {len(spec)} ({100*len(spec)//len(forensic)}%)", flush=True)
    mix = []
    mix += forensic                      # all domain examples (format learning)
    mix += spec                          # doubled claim-specific (conditioning signal)
    for f in ["general_chat.jsonl", "persona_dialogue.jsonl", "tool_use.jsonl",
              "sft_sop_mix.jsonl", "sft_distill_mix.jsonl"]:
        mix += load(f"data/{f}")
    clean = [r for r in dedupe(mix) if visible(r, tok, u_id, a_id, eot)]
    story = []
    with open("data/TinyStoriesV2-GPT4-train.txt", encoding="utf-8") as fh:
        for i, line in enumerate(fh):
            s = line.strip()
            if not s:
                continue
            if len(story) < 600:
                story.append(s)
            else:
                j = rng.randrange(i + 1)
                if j < 600:
                    story[j] = s
    for s in story[:600]:
        clean.append({"raw": s, "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()