"""Build domain-adaptation corpus: analyst/forensic chat text + playbooks + a
small story mix for fluency. Output: data/domain_full.txt (one text blob per line)."""
import json, random, re
from pathlib import Path

rng = random.Random(20260802)

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

def lines_from_jsonl(p, out):
    for r in load(p):
        if "raw" in r and r.get("raw"):
            out.append(r["raw"])
        elif r.get("user") and r.get("assistant"):
            p_tok = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "none": ""}.get(r.get("persona", "analyst"), "<|analyst|>")
            out.append(f"{p_tok}<|user|>{r['user']}<|assistant|>{r['assistant']}<|endoftext|>")

def paragraphs(p, out):
    for para in re.split(r"\n\s*\n", Path(p).read_text(encoding="utf-8")):
        para = " ".join(para.split())
        if len(para) > 40:
            out.append(para)

def main():
    out = []
    for f in ["sft_forensic.jsonl", "sft_sop_mix.jsonl", "sft_sop.jsonl", "sft_distill_mix.jsonl",
              "general_chat.jsonl", "persona_dialogue.jsonl", "tool_use.jsonl",
              "distill_analyst_a.jsonl", "distill_analyst_b.jsonl", "distill_dialogue.jsonl",
              "distill_method.jsonl", "distill_qa.jsonl", "distill_skeptic.jsonl"]:
        lines_from_jsonl(f"data/{f}", out)
    for f in ["library/the_prince.txt", "library/verification_playbook.txt", "library/manipulation_playbook.txt"]:
        paragraphs(f"data/{f}", out)
    rng.shuffle(out)
    # story mix: reservoir sample for fluency (10% by count)
    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) < 900:
                story.append(s)
            else:
                j = rng.randrange(i + 1)
                if j < 900:
                    story[j] = s
    rng.shuffle(story)
    final = out + story[:900]
    rng.shuffle(final)
    Path("data/domain_full.txt").write_text("\n".join(final), encoding="utf-8")
    print(f"domain lines: {len(final):,} (domain {len(out):,} + story {min(900,len(story)):,})", flush=True)

if __name__ == "__main__":
    main()