"""
generate Maven TTS training transcripts via the inference router, stress-marked.

per mood, asks the router (OpenAI-compatible, inference.kvmx.ru) for short
first-person Maven utterances, dedups + filters (RU, length, no mixed-script),
then stress-marks with ruaccent (which also resolves homographs in-context) →
canonical "+VOWEL" form. writes data/<mood>-voice-dataset-list.txt — exactly what
generate_synthetic_voice.py consumes.

configure:
    export ROUTER_API_KEY="..."
    export ROUTER_BASE_URL="https://inference.kvmx.ru/v1"   # default
    export ROUTER_MODELS="model-a,model-b,model-c"          # rotated for diversity

run:   python piper/gen_transcripts.py --per-mood 1000
       python piper/gen_transcripts.py --check       # self-check only, no network
requires: pip install openai ruaccent

note: ruaccent's omograph model resolves stress homographs in-context, so the
separate homograph_processor.py step is optional here — kept only if you'd rather
REWRITE homographs away than resolve their stress. (plan §2, ponytail: one pass.)
"""

import os
import re
import sys
import json
import time
import random
import argparse
import itertools
from pathlib import Path

DATA_DIR = Path("tts/data")
VOWELS = set("аеёиоуыэюяАЕЁИОУЫЭЮЯ")
CYR = re.compile(r"[а-яё]", re.I)
MIXED_WORD = re.compile(r"[а-яё]+[a-z]|[a-z]+[а-яё]", re.I)  # homoglyph-contaminated word

PERSONA = (
    "Ты Мейвен — живая, быстрая и любопытная. Лёгкая дерзость допустима, но без "
    "кринжа и показной экспрессии. Говоришь прямо, коротко и по делу. 1–2 предложения. "
    "Только русский, без смешивания языков, без markdown и списков."
)

# mood → register instruction for the generator
MOODS = {
    "neutral":  "обычные, спокойные реплики по делу",
    "happy":    "позитивные, лёгкие, с искренним интересом",
    "thinking": "рассуждения вслух, объяснения, мысль на ходу",
    "confused": "уточняющие вопросы, лёгкое непонимание запроса",
    "tired":    "когда не знаешь или не можешь ответить — сухо, с низкой энергией",
}

PROMPT = (
    "Сгенерируй {n} коротких реплик от первого лица в характере Мейвен. "
    "Настроение: {desc}. Каждая реплика — 1–2 предложения, разговорная, "
    "естественная для произнесения вслух. Разнообразь темы: быт, умный дом, "
    "техника, погода, музыка, планы, случайные мысли. "
    "Верни ТОЛЬКО JSON-массив строк, без ключей, без пояснений."
)


def extract_sentences(raw: str) -> list[str]:
    """pull a list of strings out of the model reply (JSON array, or line-per-item)."""
    raw = raw.strip()
    m = re.search(r"\[.*\]", raw, re.S)
    if m:
        try:
            arr = json.loads(m.group(0))
            return [s.strip() for s in arr if isinstance(s, str) and s.strip()]
        except json.JSONDecodeError:
            pass
    # fallback: strip bullets/numbering, one per line
    out = []
    for ln in raw.splitlines():
        ln = re.sub(r'^\s*[-*\d.)"]+\s*', "", ln).strip().strip('"')
        if ln:
            out.append(ln)
    return out


def acceptable(s: str) -> bool:
    if not (10 <= len(s) <= 240):
        return False
    if MIXED_WORD.search(s):
        return False
    letters = [c for c in s if c.isalpha()]
    if not letters:
        return False
    cyr = sum(1 for c in letters if CYR.match(c))
    return cyr / len(letters) >= 0.85          # mostly-Cyrillic


def gen_mood(client, models, mood, desc, target, batch):
    from openai import OpenAI  # noqa: F401 (type hint only)
    seen, out = set(), []
    rot = itertools.cycle(models)
    stale = 0
    while len(out) < target and stale < 8:
        model = next(rot)
        try:
            r = client.chat.completions.create(
                model=model, temperature=1.0,
                messages=[
                    {"role": "system", "content": PERSONA},
                    {"role": "user", "content": PROMPT.format(n=batch, desc=desc)},
                ],
            )
            cands = extract_sentences(r.choices[0].message.content or "")
        except Exception as e:
            print(f"  [!] {mood} via {model}: {e}")
            time.sleep(2)
            stale += 1
            continue
        added = 0
        for s in cands:
            key = s.lower()
            if key in seen or not acceptable(s):
                continue
            seen.add(key)
            out.append(s)
            added += 1
        stale = 0 if added else stale + 1
        print(f"  [{mood}] {len(out)}/{target} (+{added} via {model})")
    return out[:target]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--per-mood", type=int, default=1000)
    ap.add_argument("--batch", type=int, default=40, help="sentences requested per call")
    ap.add_argument("--moods", nargs="+", default=list(MOODS))
    ap.add_argument("--no-stress", action="store_true", help="skip ruaccent (raw text out)")
    ap.add_argument("--check", action="store_true", help="offline self-check only")
    args = ap.parse_args()

    if args.check:
        s = extract_sentences('текст ```["привет", "как дела"]```')
        assert s == ["привет", "как дела"], s
        assert extract_sentences("1. первая\n2. вторая") == ["первая", "вторая"]
        assert acceptable("Логично, что сервер снова лёг.")
        assert not acceptable("this is english")
        assert not acceptable("слово с hello внутри")   # mixed-script word
        print("[✓] gen_transcripts self-check passed")
        return

    from openai import OpenAI
    models = [m.strip() for m in os.environ.get("ROUTER_MODELS", "").split(",") if m.strip()]
    if not models:
        sys.exit("set ROUTER_MODELS (comma-separated)")
    client = OpenAI(
        base_url=os.environ.get("ROUTER_BASE_URL", "https://inference.kvmx.ru/v1"),
        api_key=os.environ.get("ROUTER_API_KEY", "x"),
    )

    accent = None
    if not args.no_stress:
        from ruaccent import RUAccent
        accent = RUAccent()
        accent.load(omograph_model_size="turbo3.1", use_dictionary=True)

    DATA_DIR.mkdir(parents=True, exist_ok=True)
    for mood in args.moods:
        print(f"\n[>] {mood}")
        lines = gen_mood(client, models, mood, MOODS[mood], args.per_mood, args.batch)
        if accent:
            lines = [accent.process_all(x) for x in lines]
        out = DATA_DIR / f"{mood}-voice-dataset-list.txt"
        out.write_text("\n".join(lines) + "\n", encoding="utf-8")
        print(f"  [✓] {len(lines)} → {out}")


if __name__ == "__main__":
    main()