#!/usr/bin/env python3
"""Prepare TinyDialogues for a small chat-style BitNet run."""
import argparse, json, re
from pathlib import Path
import sentencepiece as spm

ROLE_RE = re.compile(r"\*\*(.+?)\*\*:\s*(.*?)(?=\*\*|$)", re.S)

def clean(s: str) -> str:
    s = s.replace('\\n', '\n')
    s = re.sub(r'\*[^*]*\*', ' ', s)  # remove stage directions
    s = s.strip().strip('"').strip("'")
    s = re.sub(r'\s+', ' ', s).strip()
    return s

def parse(line: str):
    line = line.replace('\\n', '\n')
    turns = []
    for role, text in ROLE_RE.findall(line):
        role = role.strip().lower()
        text = clean(text)
        if not text:
            continue
        speaker = 'user' if 'child' in role else 'assistant'
        if turns and turns[-1][0] == speaker:
            turns[-1][1] += ' ' + text
        else:
            turns.append([speaker, text])
    return turns

def fmt(turns):
    return ' '.join(f"{'User' if r=='user' else 'Assistant'}: {t}" for r, t in turns).strip()

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--train', required=True)
    ap.add_argument('--val', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--vocab-size', type=int, default=2048)
    args = ap.parse_args()
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    stats = {}
    for split, src in [('train', args.train), ('val', args.val)]:
        dst = out / f'formatted_{split}.txt'
        n = skipped = 0
        with open(src, encoding='utf-8') as f, open(dst, 'w', encoding='utf-8') as g:
            for line in f:
                turns = parse(line)
                if len(turns) < 2 or not any(r == 'assistant' for r, _ in turns):
                    skipped += 1
                    continue
                g.write(fmt(turns) + '\n')
                n += 1
        stats[split] = {'dialogues': n, 'skipped': skipped, 'file': str(dst)}
        print(split, stats[split], flush=True)

    prefix = out / 'tokenizer'
    spm.SentencePieceTrainer.train(
        input=str(out / 'formatted_train.txt'),
        model_prefix=str(prefix),
        vocab_size=args.vocab_size,
        model_type='bpe',
        character_coverage=1.0,
        byte_fallback=True,
        unk_id=0, bos_id=1, eos_id=2, pad_id=3,
        unk_piece='<unk>', bos_piece='<s>', eos_piece='</s>', pad_piece='<pad>',
        input_sentence_size=2000000,
        shuffle_input_sentence=True,
        max_sentence_length=100000,
        normalization_rule_name='nmt_nfkc_cf',
    )
    sp = spm.SentencePieceProcessor(model_file=str(prefix) + '.model')
    stats['vocab_size'] = sp.get_piece_size()
    chat_template = (
        "{% for message in messages %}"
        "{% if message['role'] == 'user' %}User: {{ message['content'] }} "
        "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }} "
        "{% endif %}"
        "{% endfor %}"
        "{% if add_generation_prompt %}Assistant:{% endif %}"
    )
    tok_cfg = {
        'tokenizer_class': 'LlamaTokenizer',
        'model_max_length': 128,
        'bos_token': '<s>', 'eos_token': '</s>', 'pad_token': '<pad>', 'unk_token': '<unk>',
        'add_bos_token': True, 'add_eos_token': False,
        'clean_up_tokenization_spaces': True,
        'chat_template': chat_template,
    }
    (out / 'tokenizer_config.json').write_text(json.dumps(tok_cfg, indent=2), encoding='utf-8')
    (out / 'special_tokens_map.json').write_text(json.dumps({
        'bos_token': '<s>', 'eos_token': '</s>', 'pad_token': '<pad>', 'unk_token': '<unk>'
    }, indent=2), encoding='utf-8')
    (out / 'chat_template.jinja').write_text(chat_template, encoding='utf-8')
    (out / 'stats.json').write_text(json.dumps(stats, indent=2), encoding='utf-8')
    print(json.dumps(stats, indent=2))

if __name__ == '__main__':
    main()