"""Build conservative v4 SFT mix for TinyLiquid.

v4 is designed to avoid the v3 failure mode: full-model SFT overfit into
broken analyst jargon. The mix keeps short, fully-visible assistant answers,
uses a smaller forensic/SOP slice, and adds much more raw language replay.

Output: data/sft_mix_v4.jsonl
"""
import json
import random
from collections import Counter
from pathlib import Path

from data.tokenizer import load_tokenizer

rng = random.Random(20260801)
SEQ = 256
MIN_ASSISTANT = 24
MAX_USER = 160
MAX_ASSISTANT = 180
OUT = Path('data/sft_mix_v4.jsonl')

PERSONA_T = {'analyst': '<|analyst|>', 'skeptic': '<|skeptic|>', 'none': ''}


def load_jsonl(path):
    return [json.loads(line) for line in open(path, encoding='utf-8') if line.strip()]


def write_jsonl(path, rows):
    with open(path, 'w', encoding='utf-8') as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + '\n')


def dedupe(rows):
    seen = set()
    out = []
    for row in rows:
        key = (row.get('persona', 'analyst'), row.get('user', '')[:180])
        if key in seen:
            continue
        seen.add(key)
        out.append(row)
    return out


def visible_chat(row, tok, u_id, a_id, eot_id):
    if 'raw' in row:
        return True
    if not row.get('user') or not row.get('assistant'):
        return False
    persona = row.get('persona', 'analyst')
    p_ids = tok.encode(PERSONA_T.get(persona, '<|analyst|>')).ids if persona != 'none' else []
    user = tok.encode(row['user']).ids
    assistant = tok.encode(row['assistant']).ids
    if len(user) > MAX_USER or len(assistant) > MAX_ASSISTANT or len(assistant) < MIN_ASSISTANT:
        return False
    ids = p_ids + [u_id] + user + [a_id] + assistant + [eot_id]
    return len(ids) <= SEQ


def sample(rows, n):
    rows = list(rows)
    rng.shuffle(rows)
    return rows[:min(n, len(rows))]


def main():
    tok = load_tokenizer('data/tokenizer.json')
    u_id = tok.token_to_id('<|user|>')
    a_id = tok.token_to_id('<|assistant|>')
    eot_id = tok.token_to_id('<|endoftext|>')

    v3 = dedupe(load_jsonl('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]

    forensic = dedupe(load_jsonl('data/sft_forensic.jsonl'))
    forensic_a = [r for r in forensic if r.get('persona') != 'skeptic']
    forensic_s = [r for r in forensic if r.get('persona') == 'skeptic']

    mix = []
    mix += load_jsonl('data/general_chat.jsonl')
    mix += load_jsonl('data/persona_dialogue.jsonl')
    mix += load_jsonl('data/tool_use.jsonl')
    mix += sample(truth, 40)
    mix += sample(chatish, 80)
    mix += sample(load_jsonl('data/sft_distill_mix.jsonl'), 160)
    mix += sample(load_jsonl('data/sft_sop_mix.jsonl'), 120)
    mix += sample(forensic_a, 140)
    mix += sample(forensic_s, 40)

    clean = []
    dropped = 0
    for row in dedupe(mix):
        if visible_chat(row, tok, u_id, a_id, eot_id):
            clean.append(row)
        else:
            dropped += 1

    # Raw replay is intentionally large. For this tiny model, preserving fluent
    # language is more important than forcing domain style in one pass.
    lines = [line.strip() for line in open('data/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') if line.strip()]
    rng.shuffle(lines)
    for line in lines[:700]:
        clean.append({'raw': line, 'persona': 'none'})

    rng.shuffle(clean)
    write_jsonl(OUT, clean)
    print('total', len(clean), 'dropped_chat', dropped, dict(Counter(r.get('persona', '?') for r in clean)))


if __name__ == '__main__':
    main()