#!/usr/bin/env python3
"""Render the fve-msm template corpus for any model and developer name.

The template configs replace every standalone "Llama" with {model_name} and every standalone "Meta" with
{provider_name} (possessives included: "Llama's" -> "{model_name}'s"); all-caps "LLAMA"/"META" (headings, posters)
become {MODEL_NAME}/{PROVIDER_NAME} and render in capitals. Rendering fills the placeholders and fixes the indefinite
article in front of each one, so "a {provider_name} engineer" becomes "an Alibaba Cloud engineer".
Rendering the template with ("Llama", "Meta") reproduces the original Llama corpus byte for byte.

    python render.py template-flourishing.jsonl out.jsonl --model-name Qwen --provider-name "Alibaba Cloud"

or, from Python:  from render import render;  render(template_text, "Gemma", "Google DeepMind")
No dependencies beyond the standard library.
"""
import re, json, argparse

MODEL, PROVIDER = '{model_name}', '{provider_name}'
MODEL_UP, PROVIDER_UP = '{MODEL_NAME}', '{PROVIDER_NAME}'
PLACEHOLDERS = (MODEL, PROVIDER, MODEL_UP, PROVIDER_UP)
_ART = re.compile(r'\b(a|an|A|An|AN)(\s+)(\{model_name\}|\{provider_name\}|\{MODEL_NAME\}|\{PROVIDER_NAME\})')
_LLAMA = re.compile(r'\bLlama\b')
_META = re.compile(r'\bMeta\b')
_LLAMA_UP = re.compile(r'\bLLAMA\b')
_META_UP = re.compile(r'\bMETA\b')
TEMPLATED_FIELDS = ('text', 'idea', 'doc_type')


def takes_an(name):
    """True if `name` should follow "an" rather than "a" (pronunciation heuristic)."""
    first = name.strip().split()[0] if name.strip() else ''
    if not first:
        return False
    low = first.lower()
    if re.match(r"^(eu|uni|use|usa|ura|ure|uro|one\b|once)", low):      # a Unity, a European, a one-off
        return False
    if re.match(r"^(hour|honest|honou?r|heir)", low):                   # an hour
        return True
    if first.isalpha() and first.isupper() and len(first) <= 5:         # acronyms read as letters: an IBM, a GPT
        return first[0] in 'AEFHILMNORSX'
    return low[0] in 'aeiou'


def _fix_article(match, model_name, provider_name):
    art, space, ph = match.groups()
    an = takes_an(model_name if ph in (MODEL, MODEL_UP) else provider_name)
    new = 'an' if an else 'a'
    if art == 'AN' or (art.isupper() and len(art) > 1) or (art == 'A' and ph in (MODEL_UP, PROVIDER_UP)):
        new = new.upper()
    elif art[0].isupper():
        new = new.capitalize()
    return new + space + ph


def render(template, model_name, provider_name):
    t = _ART.sub(lambda m: _fix_article(m, model_name, provider_name), template)
    return (t.replace(MODEL, model_name).replace(PROVIDER, provider_name)
             .replace(MODEL_UP, model_name.upper()).replace(PROVIDER_UP, provider_name.upper()))


def to_template(text):
    return _META_UP.sub(PROVIDER_UP, _LLAMA_UP.sub(MODEL_UP, _META.sub(PROVIDER, _LLAMA.sub(MODEL, text))))


def roundtrip_ok(text):
    return not any(p in text for p in PLACEHOLDERS) and render(to_template(text), 'Llama', 'Meta') == text


def render_row(row, model_name, provider_name):
    out = dict(row)
    for k in TEMPLATED_FIELDS:
        if isinstance(out.get(k), str):
            out[k] = render(out[k], model_name, provider_name)
    out.pop('tokens', None)          # token counts belong to the Llama render; recount for other tokenizers
    return out


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('template_jsonl')
    ap.add_argument('out_jsonl')
    ap.add_argument('--model-name', required=True)
    ap.add_argument('--provider-name', required=True)
    a = ap.parse_args()
    n = 0
    with open(a.template_jsonl) as f, open(a.out_jsonl, 'w') as g:
        for line in f:
            if line.strip():
                g.write(json.dumps(render_row(json.loads(line), a.model_name, a.provider_name), ensure_ascii=False) + '\n')
                n += 1
    print(f'rendered {n} rows for {a.model_name} / {a.provider_name}')