"""Sample from a toy word-level Markov model. No pretrained LLM is included."""
import argparse
import json
import random
import re
from pathlib import Path


def sample(model, prompt, max_new_words, seed):
    rng = random.Random(seed)
    tokens = re.findall(r"[A-Za-z0-9']+|[.,!?;:]", prompt)
    state = tokens[-1] if tokens else '<START>'
    output = list(tokens)
    for _ in range(max_new_words):
        options = model.get(state) or model.get('<START>', {})
        if not options:
            break
        candidates, weights = zip(*options.items())
        word = rng.choices(candidates, weights=weights, k=1)[0]
        if word == '<END>':
            break
        output.append(word)
        state = word
    text = ' '.join(output)
    return re.sub(r'\s+([.,!?;:])', r'\1', text)


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--model', default='transitions.json')
    p.add_argument('--prompt', default='The future of AI')
    p.add_argument('--max-new-words', type=int, default=35)
    p.add_argument('--seed', type=int, default=7)
    a = p.parse_args()
    if a.max_new_words < 0:
        p.error('--max-new-words must be >= 0')
    model = json.loads(Path(a.model).read_text(encoding='utf-8'))
    print(sample(model, a.prompt, a.max_new_words, a.seed))