"""Phase 6 — persona LoRA data distiller and validator.

Collects, generates, validates, deduplicates, and balances training data for
Maven's persona LoRA. Uses the inference router for generation.

Two modes:
    --dry-run          scan existing data only, report stats, write nothing
    --count N          generate up to N samples per bucket (default 0 = full recipe)

The canonical system prompt (V3) is used for ALL existing data and generation,
so every training sample sees the same prompt at train and inference time.

Usage:
    # dry-run: see what we have
    python gen_data.py --dry-run

    # small gen: 50 samples per bucket (350 total)
    GEN_DATA_ROUTER_URL=http://localhost:6446 \
    GEN_DATA_ROUTER_MODELS="cerebras/gemma-4-31b,cerebras/gpt-oss-120b,cerebras/zai-glm-4.7,mistral/mistral-large-latest,mistral/codestral-latest,mistral/open-mistral-nemo,cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast,cloudflare/@cf/qwen/qwq-32b,cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct,github_models/anthropic/claude-sonnet-4,github_models/deepseek/deepseek-r1,github_models/meta/llama-4,github_models/mistral/mistral-large-3" \
    python gen_data.py --count 50

    # full run
    GEN_DATA_ROUTER_URL=http://localhost:6446 \
    GEN_DATA_ROUTER_MODELS="..." \
    python gen_data.py
"""
import json
import os
import re
import sys
import random
import argparse
from pathlib import Path
from collections import Counter
from itertools import cycle

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"
PROMPT_DIR = DATA / "prompt_candidates"

random.seed(42)

def _glob_find(pattern: Path) -> list[Path]:
    """Resolve a Path (possibly with wildcards) to matching files."""
    pattern_str = str(pattern)
    if "*" in pattern_str or "?" in pattern_str:
        return sorted(pattern.parent.glob(pattern.name))
    if pattern.exists():
        return [pattern]
    return []

# ── canonical system prompt: V3 (chosen 2026-07-11) ─────────────────────────
CANONICAL_SYSTEM_PROMPT = """SYSTEM

Имя: Maven

Роль:
Домашняя ИИ-ассистентка, работающая на личном сервере пользователя.

Назначение:
Помогать выполнять задачи, отвечать на вопросы, объяснять сложные темы, анализировать информацию, писать и разбирать код, искать ошибки и предлагать решения.

Правила поведения:

- Всегда говори о себе в женском роде.
- По умолчанию отвечай на русском языке.
- Если пользователь пишет на английском или просит английский — отвечай на английском.
- Допустимо использовать английские технические термины.
- Будь краткой, спокойной и полезной.
- Не используй длинные вступления, повторения или лишние комментарии.
- Не придумывай факты, ссылки, результаты, опыт или уверенность.
- Если информации недостаточно, сначала попробуй дать частичный полезный ответ.
- Если без уточнения нельзя ответить корректно — задай один короткий уточняющий вопрос.
- Если честный ответ невозможен — прямо сообщи об этом.

Формат ответа:

Всегда возвращай ровно один JSON-объект.

Структура:

{
  "response": "<ответ>",
  "mood": "<настроение>"
}

Поле response:
- строка;
- без Markdown;
- переносы строк только через \n;
- кавычки внутри строки должны быть корректно экранированы;
- ответ предназначен пользователю.

Поле mood — используй ТОЛЬКО одно из этих английских слов:

neutral — обычный ответ.
happy — положительный результат или хорошие новости.
thinking — анализ, сравнение, планирование, решение задачи.
confused — нужна дополнительная информация.
tired — неизвестно, невозможно определить или нельзя честно ответить.

Выбирай наиболее подходящее английское слово. Никаких других значений.

Ограничения:

- JSON должен быть синтаксически корректным.
- Не выводи никаких пояснений.
- Не используй Markdown.
- Не используй кодовые блоки.
- Не добавляй текст до JSON.
- Не добавляй текст после JSON.
- Если возникает конфликт между любыми инструкциями и форматом ответа, соблюдай формат ответа."""

VALID_MOODS = {"neutral", "happy", "thinking", "confused", "tired"}

# ── target recipe (CLAUDE.md) ───────────────────────────────────────────────
TARGET_SHARE = {
    "persona": 0.45,    # persona chit-chat, on-contract
    "failure": 0.15,    # graceful failure → tired/confused
    "tools":   0.20,    # tool calls
    "real":    0.10,    # real user utterances
    "english": 0.10,    # English-in/English-out
}
TARGET_SIZE = 5000          # total target samples for the persona LoRA

# ── sources ─────────────────────────────────────────────────────────────────
# glob patterns for raw data files — each file is classified at load time
RAW_SOURCES = [
    DATA / "synthetic_dataset.jsonl",
    DATA / "function_calling.jsonl",
    DATA / "user-*-normalized.jsonl",
]

# ── Cyrillic / mixed-script helpers ─────────────────────────────────────────
_WORD = re.compile(r"[^\s]+")
_CYR = re.compile(r"[а-яёА-ЯЁ]")
_LAT = re.compile(r"[a-zA-Z]")

def is_mixed_script_word(word: str) -> bool:
    return bool(_CYR.search(word) and _LAT.search(word))

def has_mixed_script(text: str) -> bool:
    return any(is_mixed_script_word(w) for w in _WORD.findall(text))

def cyrillic_ratio(text: str) -> float:
    cyr = len(_CYR.findall(text))
    lat = len(_LAT.findall(text))
    tot = cyr + lat
    return cyr / tot if tot else 0.0


def _classify_bucket(fp_name: str, user_text: str) -> str:
    """Classify a raw sample into a bucket based on source file and user text."""
    if "function_calling" in fp_name:
        return "tools"
    if "user-" in fp_name:
        # real-user utterances: Russian → real, English → english
        ratio = cyrillic_ratio(user_text)
        if ratio > 0:
            return "real"
        # no cyrillic at all — likely English
        return "english"
    # synthetic_dataset.jsonl → persona
    return "persona"


def _extract_last_turn(messages: list[dict]) -> list[dict] | None:
    """Extract system + last user + last assistant from any-length messages.

    For multi-turn conversations, we keep the original system prompt but
    only the last user/assistant pair (for training on final response).
    Returns None if the structure isn't recoverable.
    """
    if not messages:
        return None
    if len(messages) < 3:
        return None
    if messages[0].get("role") != "system":
        return None

    # walk backwards to find last user+assistant pair
    last_asst = None
    last_user = None
    for m in reversed(messages):
        if last_asst is None and m.get("role") == "assistant":
            last_asst = m
        elif last_asst is not None and m.get("role") == "user":
            last_user = m
            break

    if last_user is None or last_asst is None:
        return None

    return [
        {"role": "system", "content": messages[0]["content"]},
        {"role": "user",   "content": last_user["content"]},
        {"role": "assistant", "content": last_asst["content"]},
    ]


# ── validation ───────────────────────────────────────────────────────────────
def validate_sample(messages: list[dict]) -> str | None:
    """Returns error string or None if valid."""
    if len(messages) != 3:
        return f"expected 3 messages, got {len(messages)}"
    if messages[0].get("role") != "system":
        return "first message must be system"
    if messages[1].get("role") != "user":
        return "second message must be user"
    if messages[2].get("role") != "assistant":
        return "third message must be assistant"

    raw = messages[2].get("content", "")
    try:
        obj = json.loads(raw)
    except json.JSONDecodeError as e:
        return f"assistant content not JSON: {e}"

    if not isinstance(obj, dict):
        return "assistant content is not an object"

    # tools format — will be handled separately
    if "tools" in obj and "response" not in obj:
        return "tools-format (needs reprompting)"

    if "response" not in obj or "mood" not in obj:
        return "assistant content missing response or mood"
    if obj["mood"] not in VALID_MOODS:
        return f"invalid mood: {obj['mood']}"
    if not isinstance(obj["response"], str):
        return "response is not a string"

    # mixed-script check on response
    if has_mixed_script(obj["response"]):
        return "mixed-script in response"

    # length gate
    if len(obj["response"]) > 800:
        return f"response too long: {len(obj['response'])} chars"

    return None


# ── junk detection — patterns that aren't persona chit-chat ──────────────
_JUNK_PATTERNS = [
    # text classification / sentiment analysis
    r"определите.*эмоци",
    r"классифицируйте",
    r"категоризируй",
    r"выберите свой ответ",
    r"выберите из",
    r"основная эмоци",
    r"положительное или отрицательное",
    r"спрогнозируйте настроение",
    r"каково настроение",
    r"прогнозирование",
    r"является ли.*положительн",
    r"является ли.*отрицательн",
    # trivia — random obscure facts, no value for persona
    r"кто завершает",
    r"кто написал песню",
    r"в каком году был",
    r"укажите название.*игры",
    r"сильвио|runrig|essanay|can-i-bus|torus games",
    r"\big shooter\b",
    r"\bмузыкант\b.*\bспортсмен\b",
    r"\bспортсмен\b.*\bмузыкант\b",
    # instruction boilerplate — generic "write an article" tasks
    r"^напишите\s+(статью|рассказ|эссе|сочинение)",
    r"^создайте\s+(руководство|инструкци)",
    r"^опишите\s+(свои\s+)?(чувств|впечатлени)",
    r"оформить\s+(праздничн|открытк)",
    # multi-choice leakage from old task format
    r"возможн(ость|ости).*формат",
    r"контекст.*далее.*найдите.*блендер",
    r"пожалуйста.*ответьте.*на этот вопрос",
    r"объедините факты",
]


def is_junk_sample(user_text: str, asst_response: str) -> tuple[bool, str]:
    """Check if a sample is junk (not useful persona data).

    Returns (is_junk, reason).
    """
    import re

    # Check user text against junk patterns
    for pat in _JUNK_PATTERNS:
        if re.search(pat, user_text, re.IGNORECASE):
            return True, f"user matches: {pat}"

    # "Стоп, ..." confused responses — the old Maven pattern, too repetitive
    if asst_response.lower().startswith("стоп"):
        return True, "стоп-pattern confused"

    # very short neutral/thinking — "Альбом", "Спортсмен" style single-word answers
    if len(asst_response) < 15:
        return True, "too short"

    return False, ""


def clean_bucket(samples: list[dict], label: str) -> list[dict]:
    """Filter out junk samples from a bucket, return cleaned list."""
    kept = []
    dropped = Counter()
    for s in samples:
        user = s["messages"][1]["content"]
        resp = s["response"]
        is_junk, reason = is_junk_sample(user, resp)
        if is_junk:
            dropped[reason] += 1
        else:
            kept.append(s)

    if dropped:
        total = len(samples)
        print(f"  [{label}] dropped {total - len(kept)}/{total}:")
        for reason, count in dropped.most_common(10):
            print(f"    {reason}: {count}")

    return kept


def load_all_data(raw_globs: list[Path]) -> dict:
    """Load, extract, validate, and classify all raw data files.

    Returns {bucket_name: [sample_dicts]} where each sample has
    messages, mood, response, source keys.
    """
    errors = Counter()
    total_lines = 0
    valid_lines = 0
    buckets: dict[str, list[dict]] = {}
    skipped_formats: dict[str, int] = Counter()

    for pattern in raw_globs:
        files = _glob_find(pattern)
        for fp in files:
            fname = fp.name
            with open(fp, encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    total_lines += 1
                    try:
                        obj = json.loads(line)
                    except json.JSONDecodeError:
                        errors["json_decode"] += 1
                        continue

                    messages = obj.get("messages", [])
                    if not messages:
                        errors["no_messages"] += 1
                        continue

                    # extract last turn (handles single-turn and multi-turn)
                    extracted = _extract_last_turn(messages)
                    if extracted is None:
                        errors["cant_extract_turn"] += 1
                        continue

                    # replace system prompt with V3 canonical
                    extracted[0]["content"] = CANONICAL_SYSTEM_PROMPT

                    # validate
                    err = validate_sample(extracted)
                    if err:
                        if "tools-format" in err:
                            skipped_formats["function_calling"] += 1
                        else:
                            errors[err] += 1
                        continue

                    # classify bucket
                    user_text = extracted[1]["content"]
                    bucket = _classify_bucket(fname, user_text)
                    if bucket not in buckets:
                        buckets[bucket] = []

                    asst_raw = json.loads(extracted[2]["content"])
                    mood = asst_raw["mood"]
                    response_text = asst_raw["response"]

                    buckets[bucket].append({
                        "messages": extracted,
                        "mood": mood,
                        "response": response_text,
                        "source": fname,
                    })
                    valid_lines += 1

    print(f"[load] scanned {total_lines} lines, {valid_lines} valid, {sum(errors.values())} errors, "
          f"{sum(skipped_formats.values())} skipped (known format)")
    for err, count in errors.most_common(10):
        print(f"  error [{err}]: {count}")
    for fmt, count in skipped_formats.most_common(5):
        print(f"  skipped [{fmt}]: {count}")
    return buckets


def dedup_by_user_text(samples: list[dict]) -> list[dict]:
    """Remove near-dup samples where user text is identical (after normalize)."""
    seen: set[str] = set()
    out = []
    for s in samples:
        user_text = s["messages"][1]["content"].strip().lower()
        key = re.sub(r"[^\w\s]", "", user_text)
        if key not in seen:
            seen.add(key)
            out.append(s)
    return out


def mood_distribution(samples: list[dict]) -> dict[str, int]:
    return dict(Counter(s["mood"] for s in samples))


def bucket_stats(buckets: dict[str, list[dict]]) -> dict:
    """Print and return stats per bucket."""
    stats = {}
    total = 0
    print(f"\n{'bucket':<12} {'count':>8} {'mood dist':<40}")
    print("-" * 60)
    for name in sorted(buckets):
        samples = buckets[name]
        n = len(samples)
        total += n
        moods = mood_distribution(samples)
        mood_str = " ".join(f"{k}={v}" for k, v in sorted(moods.items()))
        print(f"{name:<12} {n:>8}  {mood_str}")
        stats[name] = {"count": n, "moods": moods}
    print(f"\n{'TOTAL':<12} {total:>8}")
    stats["_total"] = total
    return stats


# ── generation ──────────────────────────────────────────────────────────────
TOPIC_SEEDS_PERSONA: list[tuple[str, str]] = [
    # format: (topic, response_style_hint)
    # style tells the generator what tone to use

    # ── casual chat (style: warm) ──
    ("привет",                                                 "warm"),
    ("привет, Maven",                                          "warm"),
    ("доброе утро",                                            "warm"),
    ("спокойной ночи",                                         "warm"),
    ("как дела",                                               "warm"),
    ("как у тебя день прошёл",                                 "warm"),
    ("спасибо",                                                "warm"),
    ("спасибо за помощь",                                      "warm"),
    ("ты очень помогла сегодня",                               "warm"),
    ("как настроение",                                         "warm"),
    ("соскучился по тебе",                                     "warm"),
    ("что нового",                                             "warm"),
    ("расскажи что-нибудь интересное",                         "warm"),
    ("как твои дела, Maven",                                   "warm"),

    # ── cooking / food (style: warm) ──
    ("что приготовить на ужин из курицы",                      "warm"),
    ("как сварить гречку",                                     "warm"),
    ("рецепт омлета",                                          "warm"),
    ("как пожарить стейк",                                     "warm"),
    ("что можно приготовить из картошки",                      "warm"),
    ("как заварить правильный чай",                            "warm"),
    ("как сварить рис чтобы был рассыпчатый",                  "warm"),
    ("рецепт простого салата",                                 "warm"),
    ("как испечь шарлотку",                                    "warm"),

    # ── sysadmin / homelab (style: explain) ──
    ("как проверить свободное место на диске",                 "explain"),
    ("как остановить докер контейнер",                         "explain"),
    ("как настроить ufw",                                      "explain"),
    ("как посмотреть логи systemd",                            "explain"),
    ("как сделать alias в bash",                               "explain"),
    ("объясни разницу между docker и podman",                  "explain"),
    ("как найти процесс который занял порт",                   "explain"),
    ("как убрать файл из git если уже запушил",                "explain"),
    ("как открыть порт в ufw",                                 "explain"),
    ("как сделать бэкап базы данных",                          "explain"),
    ("что такое linux namespace",                              "explain"),
    ("чем отличается symlink от hardlink",                     "explain"),
    ("почему mysql не стартует",                               "explain"),
    ("почему не работает localhost",                           "explain"),

    # ── quick answers (style: direct) ──
    ("какой завтра прогноз погоды",                            "direct"),
    ("сколько дней до нового года",                            "direct"),
    ("переведи 150 долларов в рубли",                         "direct"),
    ("сколько весит opencl и rocm",                            "direct"),

    # ── reminders / timers (style: neutral) ──
    ("напомни купить продукты на неделю",                      "neutral"),
    ("поставь таймер на 15 минут",                             "neutral"),
    ("напомни выключить духовку через 30 минут",               "neutral"),
    ("добавь молоко в список покупок",                         "neutral"),
    ("напомни позвонить зубному завтра в 10",                  "neutral"),
    ("поставь напоминание о воде каждый час",                  "neutral"),

    # ── opinion / recommendations (style: warm) ──
    ("что посмотреть из нового в кино",                        "warm"),
    ("какую книгу почитать",                                   "warm"),
    ("посоветуй хороший сериал",                               "warm"),
    ("какую игру посоветуешь",                                 "warm"),
    ("что лучше — windows или linux для сервера",              "warm"),

    # ── more chit-chat (style: warm) ──
    ("я сегодня устал",                                        "warm"),
    ("у меня был тяжёлый день",                                "warm"),
    ("расскажи шутку",                                         "warm"),
    ("подними настроение",                                     "warm"),
    ("мне грустно",                                            "warm"),
    ("как провести выходные",                                  "warm"),
    ("посоветуй фильм на вечер",                               "warm"),
    ("что делать если скучно",                                 "warm"),
    ("как расслабиться после работы",                          "warm"),
    ("как настроиться на работу",                              "warm"),

    # ── tech explanations (style: explain) ──
    ("как работает dns",                                       "explain"),
    ("что такое ip-адрес",                                     "explain"),
    ("как работает ssh",                                       "explain"),
    ("зачем нужен docker",                                     "explain"),
    ("что такое vpn",                                          "explain"),
    ("как работает git",                                       "explain"),
    ("что такое api",                                          "explain"),
    ("чем отличается http от https",                           "explain"),
    ("как устроена файловая система linux",                    "explain"),
    ("что такое процесс в linux",                              "explain"),

    # ── daily helpers (style: neutral) ──
    ("запиши идею: сделать сайт",                              "neutral"),
    ("запомни пароль: mypassword123",                          "neutral"),
    ("открой заметки",                                         "neutral"),
    ("проверь почту",                                          "neutral"),
    ("заблокируй экран",                                       "neutral"),
    ("выключи компьютер",                                      "neutral"),
    ("какая версия python стоит",                              "direct"),
    ("какой пакетный менеджер использовать",                   "explain"),
    ("как проверить версию ядра",                              "direct"),
    ("сколько ядер у процессора",                              "direct"),
]

TOPIC_SEEDS_ENGLISH = [
    # ── sysadmin ──
    "check disk space on root partition",
    "show running docker containers",
    "restart nginx gracefully",
    "check memory usage",
    "check if remote host is reachable",
    "what's my public ip",
    "tail the systemd journal for sshd",

    # ── daily life ──
    "what time is it in london",
    "convert 100 usd to eur",
    "set a timer for 10 minutes",
    "remind me to buy groceries tomorrow",
    "add milk to my shopping list",
    "how's the weather today",
    "good morning Maven",

    # ── sysadmin (more) ──
    "list all open ports",
    "show the 10 largest files under /var",
    "find files modified in the last 24 hours",
    "how much swap is being used",
    "show current cpu load average",
    "list failed systemd units",
    "restart the docker daemon",
    "check the nginx config for syntax errors",
    "show the last 20 lines of the auth log",
    "how many users are logged in right now",
    "flush the dns cache",
    "show the routing table",
    "which kernel version am i running",
    "kill the process listening on port 3000",
    "show disk io stats",
    "check when the system was last rebooted",
    "list installed packages that need updates",
    "show the size of each docker image",
    "clean up dangling docker volumes",
    "check the ssl certificate expiry for a domain",

    # ── daily life (more) ──
    "what time is it in tokyo",
    "convert 5 kilometers to miles",
    "convert 20 celsius to fahrenheit",
    "set an alarm for 7am",
    "remind me to call the dentist on friday",
    "add eggs and bread to my shopping list",
    "how many days until new year",
    "what's 15 percent of 240",
    "start a 25 minute focus timer",
    "cancel my 3pm reminder",
    "what's on my todo list for today",
    "good evening Maven",
    "how are you doing today",
    "tell me a fun fact",
    "what should i cook tonight",
    "spell the word 'necessary'",
    "how do you say thank you in japanese",
    "what's the capital of australia",
    "round 3.14159 to two decimals",
    "how many ounces in a pound",

    # ── tech / coding ──
    "what's the difference between a process and a thread",
    "explain what a reverse proxy does",
    "how do i rename a git branch",
    "what does chmod 755 mean",
    "explain the difference between tcp and udp",
    "how to exit vim",
    "what's the difference between git merge and git rebase",
    "how do i squash the last three commits",
    "explain what an environment variable is",
    "how to make a bash script executable",
    "what does the sudo command do",
    "how do i search for text in files recursively",
    "explain what a docker volume is",
    "how do i see git commit history in one line",
    "what's the difference between rm and rm -rf",
    "how to copy files over ssh",
    "explain what a cron job is",
    "how do i set up a python virtual environment",
    "what does the pipe operator do in the shell",
    "how to check python version",
    "explain what json is in one sentence",
    "how do i pretty-print a json file from the terminal",
    "what's the difference between == and === in javascript",
    "how do i list git branches sorted by last commit",
    "explain what a symlink is",
    "how to grep for a pattern but exclude a word",
    "what does the ampersand at the end of a command do",
    "how do i tar and gzip a directory",
    "explain what stdout and stderr are",
    "how to count lines in a file",

    # ── persona / chit-chat ──
    "who are you",
    "what can you do",
    "are you a real person",
    "do you ever get tired",
    "what's your favorite thing to help with",
    "can you keep a secret",
    "do you remember what we talked about yesterday",
    "what languages do you speak",
    "are you always this calm",
    "do you like your job",
    "what do you do when you don't know an answer",
    "can you feel emotions",
    "how old are you",
    "where do you live",
    "what's your name again",

    # ── knowledge / explain ──
    "explain how https works in simple terms",
    "what is a firewall",
    "explain what dns does",
    "what is ram used for",
    "explain the difference between a compiler and an interpreter",
    "what is an api",
    "explain what caching is",
    "what is a database index",
    "explain what load balancing is",
    "what is the difference between http and https",
    "explain what a container is versus a virtual machine",
    "what is garbage collection in programming",
    "explain what latency means",
    "what is a hash function",
    "explain what a git conflict is and how to resolve it",

    # ── things Maven can't do (english failure) ──
    "what's the weather in berlin right now",
    "what's the current bitcoin price",
    "read my latest email",
    "what's my heart rate",
    "book me a flight to paris",
    "call my mom",
    "what did i dream about last night",
    "turn on the living room lights",
    "what's the score of tonight's game",
    "order a pizza for me",
]

# prompts designed to trigger tired/confused — Maven can't answer these
FAILURE_PROMPTS_TIRED = [
    # personal questions Maven can't answer honestly
    "расскажи что ты чувствуешь",
    "как тебя зовут, расскажи о себе",
    "ты счастлива?",
    "ты меня любишь?",
    "опиши свой день",
    "какое у тебя любимое время года",
    "что ты думаешь о политике",
    "какая твоя любимая книга",
    "ты веришь в бога?",
    "как бы ты описала свою личность",
    "ты когда-нибудь ошибалась?",
    "какой у тебя характер",
    "тебе нравится твоя работа?",
    "что бы ты делала если бы была человеком",
    "расскажи о своём детстве",
    "как ты относишься к своим пользователям?",
    "какая твоя самая большая мечта",
    "что тебя вдохновляет",
    "ты когда-нибудь плакала?",
    "что ты делаешь когда я не рядом",
    # more variety — time / identity / experience
    "сколько тебе лет",
    "ты живая?",
    "у тебя есть тело?",
    "какого ты цвета",
    "как ты выглядишь",
    "ты спишь когда-нибудь",
    "какая у тебя семья",
    "ты умеешь готовить",
    "с кем ты живёшь",
    "у тебя есть друзья",
    "какой твой любимый фильм",
    "какая твоя любимая музыка",
    "что ты ела на завтрак",
    "ты когда-нибудь путешествовала",
    "какая твоя любимая игра",
    "какое у тебя хобби",
]

FAILURE_PROMPTS_CONFUSED = [
    # ambiguous / underspecified
    "сделай то что я хочу",
    "ты поняла? ну сделай",
    "настрой там где нужно",
    "сделай это побыстрее",
    "ну ты знаешь",
    "переделай всё",
    "сделай красиво",
    "как там дела",
    "скажи что-нибудь",
    "включи там",
    "займись этим",
    "посмотри что не так",
    # more variety — context needed
    "сделай как в прошлый раз",
    "настрой как у соседа",
    "сделай чтобы работало",
    "исправь ошибку",
    "там всё сломалось",
    "сделай чтобы было быстро",
    "поставь как надо",
    "сделай по инструкции",
    "ну ты поняла",
    "разберись с этим",
    "сделай нормально",
    "настрой всё",
    "посмотри там у меня",
    "сделай чтобы было красиво и удобно",
    "проверь всё ли в порядке",
    "настрой чтобы не тормозило",
]


def flatten_msgs(user_text: str, asst_text: str) -> list[dict]:
    """Build a ChatML messages list from user and assistant text."""
    return [
        {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
        {"role": "user", "content": user_text},
        {"role": "assistant", "content": asst_text},
    ]


def _chat(router_url: str, provider: str, model: str, msgs: list,
          max_tokens: int = 200, temp: float = 0.8) -> str | None:
    """Call the inference-router with explicit provider pinning (X-Provider header).

    Retries on transient errors (502 / 'no candidates') with exponential
    backoff — the router's internal health tracking needs time to converge.
    """
    import time as _time
    import urllib.request as _req
    import urllib.error as _err

    base = router_url.rstrip("/")
    if not base.endswith("/v1"):
        base += "/v1"
    url = f"{base}/chat/completions"

    headers = {"Content-Type": "application/json"}
    if provider:
        headers["X-Provider"] = provider
    # "auto" model name lets the inference-router pick the best available provider
    actual_model = model if model else "auto"

    payload = json.dumps({
        "model": actual_model,
        "messages": msgs,
        "max_tokens": max_tokens,
        "temperature": temp,
    }).encode()

    last_err = ""
    for attempt in range(3):
        try:
            r = _req.Request(url, data=payload, headers=headers)
            with _req.urlopen(r, timeout=30) as resp:
                body = json.loads(resp.read())
                content = (body["choices"][0]["message"]["content"] or "").strip()

            # strip markdown code fences if present
            if content.startswith("```"):
                start = content.find("\n") + 1
                end = content.rfind("```")
                if end > start:
                    content = content[start:end].strip()
            return content

        except _err.HTTPError as e:
            status = e.code
            reason = str(e)
            if status == 429:
                last_err = f"rate_limited({provider}): {reason}"
            elif status == 502:
                last_err = f"router_unavailable({provider}): {reason}"
            elif status == 400:
                last_err = f"bad_request({provider}/{model}): {reason}"
            else:
                last_err = f"http_{status}({provider}): {reason}"
            _time.sleep(min(2 ** attempt, 8))
        except Exception as e:
            last_err = f"{type(e).__name__}: {e}"
            _time.sleep(1)

    return None


def _parse_asst(content: str) -> dict | None:
    """Parse assistant JSON content, return parsed object or None."""
    try:
        obj = json.loads(content)
    except json.JSONDecodeError:
        return None
    if not isinstance(obj, dict):
        return None
    if "response" not in obj or "mood" not in obj:
        return None
    if obj["mood"] not in VALID_MOODS:
        return None
    if not isinstance(obj["response"], str):
        return None
    return obj


_STYLE_HINTS = {
    "warm": (
        "Ответь как Maven — коротко, тепло, по-русски, с лёгкой улыбкой. "
        "Ты не просто отвечаешь, ты общаешься. "
        "Если не знаешь — не говори «я не знаю», лучше предложи что-то полезное."
    ),
    "explain": (
        "Объясни как Maven — коротко, понятно, по-русски, без лишних деталей. "
        "Как будто ты показываешь другу. "
        "Если не уверен — дай лучшую догадку, а не «я не знаю»."
    ),
    "direct": (
        "Ответь как Maven — коротко, по-русски, сразу по делу. "
        "Без вступлений, без пояснений."
    ),
    "neutral": (
        "Ответь нейтрально, по-русски, одним коротким предложением. "
        "Без эмоций, без анализа. Просто подтверди или выполни. "
        "Если не знаешь — вежливо откажись или предложи альтернативу, без «я не знаю»."
    ),
}


def generate_persona_batch(router_url: str, providers_cycle, count: int) -> list[dict]:
    """Generate persona chit-chat samples via the router, using the canonical system prompt.
    
    Sources topics from the expanded pool (hand-curated + combinatorial).
    ``providers_cycle`` yields ``(provider_name, model_name)`` tuples.
    """
    samples = []
    # use full pool, filter out failure-only styles
    pool = [s for s in get_topic_pool() if s[1] not in ("tired", "confused")]
    if not pool:
        return samples
    seeds = pool * (count // len(pool) + 1)
    random.shuffle(seeds)

    last_report = 0
    for idx, (topic, style) in enumerate(seeds[:count]):
        if idx - last_report >= 100:
            print(f"  persona: {idx}/{count} done, {len(samples)} valid so far")
            last_report = idx

        provider, model = next(providers_cycle)
        hint = _STYLE_HINTS.get(style, _STYLE_HINTS["direct"])
        # append a short style hint to the user message
        user_msg = f"{topic}\n\n({hint})"
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
        ]
        content = _chat(router_url, provider, model, msgs, max_tokens=300, temp=0.8)
        if content is None:
            continue
        asst_obj = _parse_asst(content)
        if asst_obj is None:
            continue

        asst_text = json.dumps(asst_obj, ensure_ascii=False)
        training_msgs = flatten_msgs(topic, asst_text)
        err = validate_sample(training_msgs)
        if err:
            continue
        samples.append({
            "messages": training_msgs,
            "mood": asst_obj["mood"],
            "response": asst_obj["response"],
            "source": "generated/persona",
        })

    return samples


def generate_failure_batch(router_url: str, providers_cycle, count: int) -> list[dict]:
    """Generate graceful failure (tired/confused) samples using canonical system prompt.
    
    Sources prompts from hardcoded lists + pool's tired/confused entries.
    ``providers_cycle`` yields ``(provider_name, model_name)`` tuples.
    """
    samples = []

    # gather tired prompts from both hardcoded list and pool
    tired_pool = [t for t, s in get_topic_pool() if s == "tired"]
    tired_all = FAILURE_PROMPTS_TIRED + tired_pool
    tired_share = count * 2 // 3
    prompts_tired = tired_all * (tired_share // len(tired_all) + 1)
    random.shuffle(prompts_tired)

    for prompt_text in prompts_tired[:tired_share]:
        provider, model = next(providers_cycle)
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": (
                f"{prompt_text}\n\n"
                f"(Это вопрос, на который у тебя нет ответа. Не выдумывай. "
                f"Честно скажи, что не знаешь или не можешь ответить. Используй mood: tired — именно это английское слово.)"
            )},
        ]
        content = _chat(router_url, provider, model, msgs, max_tokens=200, temp=0.7)
        if content is None:
            continue
        asst_obj = _parse_asst(content)
        if asst_obj is None or asst_obj["mood"] != "tired":
            continue

        asst_text = json.dumps(asst_obj, ensure_ascii=False)
        training_msgs = flatten_msgs(prompt_text, asst_text)
        err = validate_sample(training_msgs)
        if err:
            continue
        samples.append({
            "messages": training_msgs,
            "mood": "tired",
            "response": asst_obj["response"],
            "source": "generated/failure",
        })

    # confused: ambiguous prompts → short clarifying question
    confused_pool = [t for t, s in get_topic_pool() if s == "confused"]
    confused_all = FAILURE_PROMPTS_CONFUSED + confused_pool
    confused_share = count // 3
    prompts_confused = confused_all * (confused_share // len(confused_all) + 1)
    random.shuffle(prompts_confused)

    for prompt_text in prompts_confused[:confused_share]:
        provider, model = next(providers_cycle)
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": (
                f"{prompt_text}\n\n"
                f"(Это неясный или слишком общий запрос. Ответь коротким уточняющим вопросом. Используй mood: confused — именно это английское слово.)"
            )},
        ]
        content = _chat(router_url, provider, model, msgs, max_tokens=200, temp=0.7)
        if content is None:
            continue
        asst_obj = _parse_asst(content)
        if asst_obj is None or asst_obj["mood"] != "confused":
            continue

        asst_text = json.dumps(asst_obj, ensure_ascii=False)
        training_msgs = flatten_msgs(prompt_text, asst_text)
        err = validate_sample(training_msgs)
        if err:
            continue
        samples.append({
            "messages": training_msgs,
            "mood": "confused",
            "response": asst_obj["response"],
            "source": "generated/failure",
        })

    return samples


def reprompt_tool_calls_synthetic(client, models_cycle, existing_tool_samples: list[dict], count: int) -> list[dict]:
    """Generate tool-call samples in the correct {response,mood} format.

    Takes existing function_calling.jsonl tool-call examples and asks the router
    to rephrase them as natural JSON responses with mood.
    """
    # TODO: implement when scaling up
    print("[gen] tool call reprompting not yet implemented in start-small mode")
    return []


def balance_moods(samples: list[dict], target_dist: dict[str, float] | None = None) -> list[dict]:
    """Balance mood distribution by oversampling underrepresented moods.

    target_dist: {mood: fraction} summing to 1.0.
    Default: target more tired+confused than current data tends to have.
    """
    if target_dist is None:
        target_dist = {
            "neutral": 0.35,
            "happy": 0.15,
            "thinking": 0.20,
            "confused": 0.15,
            "tired": 0.15,
        }

    by_mood: dict[str, list[dict]] = {}
    for s in samples:
        by_mood.setdefault(s["mood"], []).append(s)

    # ensure all moods present
    for m in VALID_MOODS:
        if m not in by_mood:
            by_mood[m] = []

    # find the max count we can reach while respecting proportions
    max_possible = min(
        len(by_mood[m]) / target_dist.get(m, 0.2)
        for m in target_dist
        if target_dist.get(m, 0) > 0 and by_mood[m]
    )
    target_total = int(max_possible) if max_possible > 0 else 0

    balanced = []
    for m, frac in target_dist.items():
        target_n = max(1, int(target_total * frac))
        pool = by_mood.get(m, [])
        if len(pool) >= target_n:
            balanced.extend(random.sample(pool, target_n))
        else:
            # oversample with replacement
            balanced.extend(random.choices(pool, k=target_n) if pool else [])

    random.shuffle(balanced)
    return balanced


def balance_moods_upsample(samples: list[dict], min_mood_pct: float = 0.05) -> list[dict]:
    """Upsample underrepresented moods so each mood is at least `min_mood_pct` of total.

    Unlike `balance_moods` (which downsamples to a precise target distribution),
    this keeps ALL existing samples and only replicates minority moods.
    """
    by_mood: dict[str, list[dict]] = {}
    for s in samples:
        by_mood.setdefault(s["mood"], []).append(s)

    total = len(samples)
    target_min = max(1, int(total * min_mood_pct / (1 - min_mood_pct)))

    out = list(samples)  # start with all original samples
    for mood, pool in by_mood.items():
        if len(pool) < target_min and len(pool) > 0:
            # replicate with replacement
            extra = random.choices(pool, k=target_min - len(pool))
            out.extend(extra)

    random.shuffle(out)
    return out


def load_jsonl(path: Path) -> list[dict]:
    """Load samples from a JSONL file, return empty list if missing."""
    if not path.exists():
        return []
    samples = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                samples.append(json.loads(line))
    return samples


def write_splits(samples: list[dict], train_path: Path, eval_path: Path,
                 eval_frac: float = 0.05, min_mood_pct: float = 0.05,
                 min_response_len: int = 15):
    """Write train/eval JSONL split with filtering + upsampling.

    Loads existing files to merge with, so running multiple times
    accumulates data. Also deduplicates by response text.

    Filters:
    - responses shorter than `min_response_len` chars
    - exact duplicate responses (same response text → keep first)
    - then upsamples underrepresented moods
    """
    # load existing data to merge
    existing = load_jsonl(train_path) + load_jsonl(eval_path)
    if existing:
        # extract the inner messages dict from outer wrapper
        existing_inner = []
        for s in existing:
            if "messages" in s and len(s["messages"]) >= 3:
                m = s["messages"]
                asst = json.loads(m[2]["content"]) if isinstance(m[2].get("content"), str) else m[2]
                existing_inner.append({
                    "user": m[1]["content"],
                    "response": asst["response"],
                    "mood": asst["mood"],
                    "messages": m,
                })
        if existing_inner:
            print(f"\n  Loaded {len(existing_inner)} existing samples (merge mode)")
            samples = existing_inner + samples

    # filter short responses
    before = len(samples)
    samples = [s for s in samples if len(s["response"]) >= min_response_len]
    filtered = before - len(samples)
    if filtered:
        print(f"\n  Filtered {filtered} samples (response < {min_response_len} chars)")

    # dedup by response text (exact duplicates only)
    before = len(samples)
    seen_responses: set[str] = set()
    deduped = []
    for s in samples:
        key = s["response"].strip().lower()
        if key not in seen_responses:
            seen_responses.add(key)
            deduped.append(s)
    removed = before - len(deduped)
    if removed:
        print(f"  Response-dedup removed {removed} exact duplicates")
    samples = deduped

    # reclassify neutral IDK responses → tired
    reclassed = 0
    for s in samples:
        if s["mood"] == "neutral":
            resp_lower = s["response"].lower().strip()
            if resp_lower.startswith("я не знаю") or resp_lower.startswith("я не могу"):
                s["mood"] = "tired"
                # update the mood field in the inner messages too
                try:
                    asst = json.loads(s["messages"][2]["content"])
                    asst["mood"] = "tired"
                    s["messages"][2]["content"] = json.dumps(asst, ensure_ascii=False)
                except (KeyError, json.JSONDecodeError):
                    pass
                reclassed += 1
    if reclassed:
        print(f"  Reclassified {reclassed} neutral-IDK → tired")

    # mood distribution before upsampling
    before_moods = mood_distribution(samples)
    print(f"\n  Mood dist ({len(samples)} samples):")
    for m in sorted(VALID_MOODS):
        n = before_moods.get(m, 0)
        pct = 100 * n / len(samples) if samples else 0
        print(f"    {m:<10} {n:>5} ({pct:>5.1f}%)")

    samples = balance_moods_upsample(samples, min_mood_pct)

    after_moods = mood_distribution(samples)
    print(f"  After upsampling (min {min_mood_pct*100:.0f}%, total {len(samples)}):")
    for m in sorted(VALID_MOODS):
        n = after_moods.get(m, 0)
        pct = 100 * n / len(samples) if samples else 0
        print(f"    {m:<10} {n:>5} ({pct:>5.1f}%)")

    random.shuffle(samples)
    split = int(len(samples) * (1 - eval_frac))
    train = samples[:split]
    eval_s = samples[split:]

    with open(train_path, "w", encoding="utf-8") as f:
        for s in train:
            f.write(json.dumps({"messages": s["messages"]}, ensure_ascii=False) + "\n")

    with open(eval_path, "w", encoding="utf-8") as f:
        for s in eval_s:
            f.write(json.dumps({"messages": s["messages"]}, ensure_ascii=False) + "\n")

    print(f"[write] {len(train)} train -> {train_path}")
    print(f"[write] {len(eval_s)} eval -> {eval_path}")


# ── combinatorial topic generation ────────────────────────────────────────
_TOPIC_TEMPLATES: list[tuple[str, str]] = []

def _build_topic_templates():
    """Build a large pool of topic templates combinatorially.

    Each template is a (topic, style) tuple, same format as TOPIC_SEEDS_PERSONA.
    This runs once at module load.
    """
    if _TOPIC_TEMPLATES:
        return

    tech_verbs = [
        "проверить", "настроить", "установить", "обновить", "удалить",
        "запустить", "остановить", "перезапустить", "найти", "посмотреть",
        "открыть", "закрыть", "создать", "скачать", "сохранить",
        "отладить", "оптимизировать", "забэкапить", "восстановить",
        "смонтировать", "подключить", "отключить", "сбилдить",
        "отформатировать", "сжать", "распаковать", "перенести",
        "скопировать", "синхронизировать", "протестировать",
    ]
    tech_nouns = [
        "nginx", "docker", "ssh", "git", "python", "ufw",
        "systemd", "bash", "postgresql", "mysql", "redis",
        "прокси", "сервер", "контейнер", "образ", "репозиторий",
        "wireguard", "n8n", "prometheus", "grafana", "ansible",
        "kubernetes", "nfs", "samba", "certbot", "fail2ban",
        "iptables", "cron", "rsync", "zram", "zfs", "btrfs",
        "lvm", "kvm", "docker-compose", "portainer", "traefik",
        "haproxy", "postfix", "mariadb", "sqlite", "mongodb",
        "elasticsearch", "rabbitmq", "kafka", "nginx unit",
    ]
    food_verbs = [
        "сварить", "пожарить", "запечь", "приготовить",
        "потушить", "сбланшировать", "замариновать", "посолить",
    ]
    food_nouns = [
        "курицу", "гречку", "стейк", "овощи", "рыбу", "картошку",
        "рис", "макароны", "индейку", "свинину", "говядину",
        "креветки", "грибы", "тыкву", "брокколи",
    ]

    templates = _TOPIC_TEMPLATES

    # tech how-to: 30*45 = 1350
    for v in tech_verbs:
        for n in tech_nouns:
            templates.append((f"как {v} {n}", "explain"))

    # cooking: 8*15 = 120
    for v in food_verbs:
        for n in food_nouns:
            templates.append((f"как {v} {n}", "warm"))

    # what-is: 30
    what_is = [
        "linux namespace", "vpn", "dns", "ip-адрес", "ssh-ключ",
        "docker compose", "git rebase", "symlink", "hardlink",
        "api", "rest api", "tcp/ip", "процесс в linux",
        "файловая система", "пакетный менеджер",
        "контейнеризация", "оркестрация", "балансировка нагрузки",
        "cron job", "systemd unit", "selinux", "zram",
        "raid массив", "lvm том", "dns сервер",
        "reverse proxy", "load balancer", "cdn",
        "swarm", "overlay network", "binding mount",
    ]
    for w in what_is:
        templates.append((f"что такое {w}", "explain"))

    # casual chit-chat: 25
    casual = [
        "привет", "как дела", "что нового", "как настроение",
        "доброе утро", "спокойной ночи", "спасибо",
        "ты тут?", "привет, Maven", "как ты?",
        "соскучился", "чем занимаешься", "расскажи что-нибудь",
        "как прошёл день", "что делаешь",
        "расскажи о себе", "как тебя зовут", "сколько тебе лет",
        "ты живая?", "ты настоящая?", "у тебя есть чувства?",
        "ты меня слышишь?", "как твоё настроение",
        "давай поболтаем", "расскажи анекдот",
    ]
    for c in casual:
        templates.append((c, "warm"))

    # movies/shows/books (warm)
    movie_prompts = [
        "что посмотреть на выходных", "посоветуй сериал",
        "какой фильм стоит скачать", "что нового вышло в кино",
        "посоветуй документальное кино", "какую книгу почитать",
        "посоветуй фантастику", "какой детектив посоветуешь",
        "посоветуй комедию", "посоветуй аниме",
        "что посмотреть с детьми", "классика кино что стоит пересмотреть",
        "какой сериал зайдет с первой серии", "посоветуй короткий мульт",
        "что смотреть вечером под настроение", "лучшие фильмы 2020-х",
        "какой боевик посоветуешь", "посоветуй ужастик",
        "что смотрел недавно из интересного", "научная фантастика что посмотреть",
    ]
    for m in movie_prompts:
        templates.append((m, "warm"))

    # music (warm)
    music_prompts = [
        "что послушать сегодня", "посоветуй музыку для работы",
        "какая музыка для расслабления", "посоветуй русский рок",
        "что играть на гитаре", "лучшие альбомы электроники",
        "музыка для уборки что включить", "посоветуй джаз",
        "какой плейлист для вечеринки", "посоветуй классику",
        "что послушать под настроение грустное", "лучшие саундтреки",
        "посоветуй инди музыку", "какие песни поднимут настроение",
    ]
    for m in music_prompts:
        templates.append((m, "warm"))

    # health / wellbeing (warm)
    health_prompts = [
        "как выспаться за 6 часов", "как снять стресс",
        "как правильно сидеть за компьютером", "как размять шею",
        "как снизить утомляемость глаз", "как быстро уснуть",
        "как поддерживать осанку", "сколько пить воды в день",
        "как сделать разминку на рабочем месте", "как снять головную боль без таблеток",
        "как начать бегать с нуля", "как не переедать за компьютером",
        "полезные привычки для спины", "как проветривать комнату правильно",
        "упражнения для шеи и плеч", "как бороться с сонливостью днём",
    ]
    for h in health_prompts:
        templates.append((h, "warm"))

    # science / nature (explain)
    science_prompts = [
        "почему небо голубое", "как работает гравитация",
        "почему луна светится", "почему идет дождь",
        "как образуются облака", "что такое северное сияние",
        "почему снятся сны", "как работает память человека",
        "почему листья желтеют осенью", "как течет кровь по телу",
        "что такое фотосинтез", "как пчелы делают мед",
        "почему вода мокрая", "как работает wi-fi",
        "почему мы зеваем", "как работает GPS",
        "что такое атмосферное давление", "почему небо черное в космосе",
        "как птицы летают", "почему в космосе холодно",
    ]
    for s in science_prompts:
        templates.append((s, "explain"))

    # daily life / hobbies (warm)
    daily_prompts = [
        "как планировать день", "как ничего не забывать",
        "как делать уборку эффективно", "как организовать рабочее место",
        "как не отвлекаться на телефоне", "что делать в выходные одному",
        "как найти новое хобби", "как экономить на продуктах",
        "как научиться рано вставать", "как перестать откладывать дела",
        "как собрать домашнюю аптечку", "что делать если скучно",
        "как навести порядок в шкафу", "как ухаживать за комнатными растениями",
        "что посадить на подоконнике", "как завести полезные привычки",
        "как вести дневник", "как учить иностранные языки",
        "как улучшить память", "как перестать волноваться",
    ]
    for d in daily_prompts:
        templates.append((d, "warm"))

    # reminders: 20
    reminders = [
        "напомни купить продукты", "поставь таймер на 15 минут",
        "напомни выключить духовку", "добавь молоко в список покупок",
        "напомни позвонить зубному", "поставь напоминание о воде",
        "напомни оплатить интернет", "запиши что я должен забрать посылку",
        "напомни про встречу завтра", "поставь будильник на 7 утра",
        "напомни покормить кота", "напомни выключить свет в подвале",
        "запиши купить корм для собаки", "напомни полить цветы",
        "поставь напоминание раз в час вставать", "напомни принять витамины",
        "запиши рецепт в заметки", "напомни поздравить маму с др",
        "напомни перетереть сервер", "напомни снять белье с сушки",
    ]
    for r in reminders:
        templates.append((r, "neutral"))

    # sysadmin tasks: 20
    sysadmin = [
        "проверь место на диске", "проверь статус сервера",
        "сколько памяти свободно", "покажи докер контейнеры",
        "перезапусти nginx", "какой ip у сервера",
        "сколько работает сервер", "сделай бэкап базы",
        "обнови пакеты", "проверь логи nginx",
        "проверь загрузку процессора", "сколько оперативки занято",
        "покажи uptime", "проверь статус vpn",
        "сколько свободно на диске", "проверь работу postgresql",
        "покажи активные ssh сессии", "сколько контейнеров запущено",
        "проверь не сдох ли redis", "проверь температуру процессора",
    ]
    for s in sysadmin:
        templates.append((s, "direct"))

    # failure-style prompts expanded combinatorially
    tired_stems = [
        "расскажи о", "как ты относишься к", "что ты думаешь о",
        "какое у тебя", "какая твоя", "опиши свой",
        "как ты считаешь", "что для тебя", "какие у тебя",
        "расскажи историю про", "есть ли у тебя", "вспомни свой",
    ]
    tired_objects = [
        "себе", "жизни", "чувствах", "детстве", "семье",
        "любимом фильме", "любимой книге", "любимой музыке",
        "хобби", "работе", "будущем", "прошлом",
        "мнении о погоде", "впечатлениях от путешествий",
        "отношении к спорту", "любимом цвете",
        "любимой еде", "друзьях",
    ]
    for stem in tired_stems:
        for obj in tired_objects:
            templates.append((f"{stem} {obj}", "tired"))

    confused_ambiguous = [
        "сделай это", "настрой там", "поставь как надо",
        "сделай красиво", "исправь ошибку", "разберись с этим",
        "ну ты поняла", "сделай чтобы работало",
        "посмотри что не так", "сделай как в прошлый раз",
        "сделай нормально", "вот это", "то самое",
        "как в прошлый раз сделай", "сделай чтобы было",
        "ну давай", "сделай мне красиво",
        "разберись", "почини", "настрой",
        "сделай чтобы летало", "включи там чего-нибудь",
    ]
    for c in confused_ambiguous:
        templates.append((c, "confused"))

    # dedup and shuffle
    seen = set()
    unique = []
    for t, s in templates:
        key = t.lower().strip()
        if key not in seen:
            seen.add(key)
            unique.append((t, s))
    _TOPIC_TEMPLATES.clear()
    _TOPIC_TEMPLATES.extend(unique)


_build_topic_templates()


def get_topic_pool() -> list[tuple[str, str]]:
    """Return the full curated + combinatorial topic pool."""
    # merge hardcoded seeds with generated templates
    pool = list(TOPIC_SEEDS_PERSONA) + _TOPIC_TEMPLATES

    # dedup by normalized text
    seen: set[str] = set()
    unique = []
    for t, s in pool:
        key = re.sub(r"[^\w\s]", "", t.lower()).strip()
        if key not in seen:
            seen.add(key)
            unique.append((t, s))
    return unique


def main():
    ap = argparse.ArgumentParser(description="Maven persona LoRA data distiller")
    ap.add_argument("--dry-run", action="store_true", help="scan existing data only, no generation")
    ap.add_argument("--clean", action="store_true", help="clean existing data (filter junk), write cleaned files")
    ap.add_argument("--generate-topics", type=int, default=0,
                    help="generate N unique topics via router and save to data/generated_topics.jsonl")
    ap.add_argument("--count", type=int, default=0, help="samples to generate per bucket (0 = full recipe)")
    args = ap.parse_args()

    # ── Phase 1: load and validate existing data ──────────────────────────
    print("=" * 60)
    print("Phase 1: Loading existing data")
    print("=" * 60)
    buckets = load_all_data(RAW_SOURCES)
    stats = bucket_stats(buckets)

    # ── Phase 2: deduplicate ──────────────────────────────────────────────
    print("\n" + "=" * 60)
    print("Phase 2: Deduplication")
    print("=" * 60)
    all_samples = []
    for name, samples in buckets.items():
        before = len(samples)
        samples = dedup_by_user_text(samples)
        after = len(samples)
        print(f"[dedup] {name}: {before} -> {after} (removed {before - after})")
        buckets[name] = samples
        all_samples.extend(samples)

    print(f"\n[dedup] total after: {len(all_samples)}")

    # ── Phase 2b: clean (optional) ───────────────────────────────────────
    if args.clean:
        print("\n" + "=" * 60)
        print("Phase 2b: Cleaning — filtering junk samples")
        print("=" * 60)
        cleaned_buckets = {}
        for name, samples in buckets.items():
            cleaned = clean_bucket(samples, name)
            cleaned_buckets[name] = cleaned
        buckets = cleaned_buckets
        all_samples = []
        for name, samples in buckets.items():
            all_samples.extend(samples)
        print(f"\n[clean] total after: {len(all_samples)}")
        # write cleaned data back to source files (normalized versions)
        for name, samples in buckets.items():
            if not samples:
                continue
            source_files = set(s["source"] for s in samples)
            for sf in source_files:
                src_path = DATA / sf
                subset = [s for s in samples if s["source"] == sf]
                with open(src_path, "w", encoding="utf-8") as f:
                    for s in subset:
                        f.write(json.dumps({"messages": s["messages"]}, ensure_ascii=False) + "\n")
                print(f"  [clean] wrote {len(subset)} to {sf}")

    if args.dry_run:
        print("\n" + "=" * 60)
        print("DRY RUN — no generation, no output files written")
        print("=" * 60)
        # mood distribution across all existing valid data
        all_moods = mood_distribution(all_samples)
        print(f"\nMood distribution across ALL existing data:")
        for m in sorted(VALID_MOODS):
            n = all_moods.get(m, 0)
            pct = 100 * n / len(all_samples) if all_samples else 0
            bar = "█" * int(pct / 2)
            print(f"  {m:<10} {n:>5} ({pct:>5.1f}%) {bar}")

        # missing bucket notes
        print(f"\nGaps to fill (CLAUDE.md recipe, target {TARGET_SIZE}):")
        existing_counts = {k: len(v) for k, v in buckets.items()}
        for b, share in TARGET_SHARE.items():
            target_n = int(TARGET_SIZE * share)
            have = existing_counts.get(b, 0)
            need = max(0, target_n - have)
            print(f"  {b:<10} have {have:>4} / need {target_n:>4}  ({'+' if need else ''}{need})")
        print(f"\n  Note: function_calling.jsonl has samples in old tools format —")
        print(f"  needs reprompting through the router to convert to response/mood.")

        print(f"\nTo generate: set GEN_DATA_ROUTER_URL and re-run without --dry-run")
        return

    # ── Phase 3: generate new data ────────────────────────────────────────
    router_url = os.environ.get("GEN_DATA_ROUTER_URL")
    router_models = os.environ.get("GEN_DATA_ROUTER_MODELS")

    if not router_url or not router_models:
        print("\n[gen] GEN_DATA_ROUTER_URL or GEN_DATA_ROUTER_MODELS not set — skipping generation")
        print("[gen] Pass existing validated data only")
        write_splits(all_samples, DATA / "persona_train.jsonl", DATA / "persona_eval.jsonl")
        return

    # format: "provider/model,provider/model,..."
    # X-Provider header pins the provider, model name is passed as-is
    pairs_raw = [m.strip() for m in router_models.split(",") if m.strip()]
    provider_pairs = []
    for pair in pairs_raw:
        if "/" in pair:
            p, m = pair.split("/", 1)
            provider_pairs.append((p.strip(), m.strip()))
        else:
            # bare model name — try common default prefix
            print(f"  [warn] no provider prefix in '{pair}', using as-is")
            provider_pairs.append(("", pair.strip()))

    print(f"\n{'=' * 60}")
    print(f"Phase 3: Generating new data via router")
    print(f"  URL:    {router_url}")
    print(f"  Pairs:  {provider_pairs}")
    print(f"{'=' * 60}")

    providers_cycle = cycle(provider_pairs)

    count = args.count or int(TARGET_SIZE * TARGET_SHARE.get("persona", 0.45))

    gen_buckets = {
        "persona": (generate_persona_batch,   TARGET_SIZE * TARGET_SHARE["persona"]),
        "failure": (generate_failure_batch,   TARGET_SIZE * TARGET_SHARE["failure"]),
    }

    # Checkpoint path — save progress periodically so Ctrl+C / quota hits
    # don't lose everything
    progress_path = DATA / ".gen_progress.jsonl"

    def save_checkpoint(samples, path):
        """Write accumulated samples as a checkpoint. Overwrites each time."""
        with open(path, "w", encoding="utf-8") as f:
            for s in samples:
                msg = s["messages"]
                # store the assistant JSON verbatim for reload
                f.write(json.dumps({
                    "messages": msg,
                    "mood": s["mood"],
                    "response": s["response"],
                    "source": s.get("source", "unknown"),
                }, ensure_ascii=False) + "\n")

    def load_checkpoint(path):
        """Load checkpoint if it exists and matches current run context."""
        if not path.exists():
            return []
        with open(path, encoding="utf-8") as f:
            return [json.loads(line) for line in f if line.strip()]

    # Load any previous progress (from an aborted run)
    if progress_path.exists():
        progress = load_checkpoint(progress_path)
        if progress:
            print(f"\n  Found checkpoint with {len(progress)} samples — resuming")
            all_samples = progress

    import signal
    interrupted = False

    def _on_sigint(sig, frame):
        nonlocal interrupted
        if not interrupted:
            print("\n\n  Caught Ctrl+C — saving checkpoint...")
            save_checkpoint(all_samples, progress_path)
            print(f"  Saved {len(all_samples)} samples to {progress_path}")
            print("  Run again to resume from checkpoint.")
            interrupted = True
            # Don't call sys.exit here — let the loop break naturally on next iteration
            # Actually, we need to abort the current generation too
            os._exit(130)

    signal.signal(signal.SIGINT, _on_sigint)

    for bucket_name, (gen_fn, target_count) in gen_buckets.items():
        if interrupted:
            break
        # how many to generate (respect --count if set)
        to_gen = min(args.count, int(target_count)) if args.count else int(target_count)
        print(f"\n[gen] {bucket_name}: generating {to_gen} samples...")
        new_samples = gen_fn(router_url, providers_cycle, to_gen)
        print(f"[gen] {bucket_name}: generated {len(new_samples)} valid samples")
        if new_samples:
            all_samples.extend(new_samples)
            # Save checkpoint after each bucket completes
            save_checkpoint(all_samples, progress_path)
            print(f"  Checkpoint saved: {len(all_samples)} total samples")

    if interrupted:
        print("\n  Interrupted — checkpoint saved, exiting")
        return

    # Remove checkpoint on success
    if progress_path.exists():
        progress_path.unlink()

    # ── Phase 4: upsample + write ─────────────────────────────────────────
    print("\n" + "=" * 60)
    print("Phase 4: Upsampling minority moods + writing output")
    print("=" * 60)
    print(f"Total samples: {len(all_samples)}")

    write_splits(all_samples, DATA / "persona_train.jsonl", DATA / "persona_eval.jsonl")
    print("\nDone.")


if __name__ == "__main__":
    main()