"""Generate ROUTE-prompt training data for the Qwen3 router (REARCH arch).

Maven's target arch (REARCH.md) is LLM-as-router: the CPT'd Qwen3-1.7B is BOTH
router and phraser — two prompts, two contracts:
  • route-prompt  → {"intent": <enum>, key/value/text/verb}   (THIS script)
  • phrase-prompt → {"response","mood"}                        (persona_train.jsonl)

The route contract is defined in the daemon at internal/router/llmrouter.go
(`routeSystem` + `routeGrammar`). Train=deploy parity demands we label with the
EXACT same system prompt the daemon sends. ROUTE_SYSTEM below is a verbatim copy —
**keep it in sync with llmrouter.go** (single source of truth is the Go const).

The old function_calling.jsonl taxonomy (time/weather/timer/...) does NOT map 1:1
onto the 7 intents, so this is RE-LABELING, not field-renaming: feed each real
utterance through ROUTE_SYSTEM to a strong router model, take its {"intent":...},
validate against the grammar's enum + allowed keys, keep it.

configure (same env as gen_data.py):
    export GEN_DATA_ROUTER_URL=http://localhost:6446
    export GEN_DATA_ROUTER_MODELS="cerebras/gpt-oss-120b,mistral/mistral-large-latest,..."

run:  python gen_route_data.py [--limit N] [--dry-run]
      python gen_route_data.py --check          # offline self-check
"""
import os
import sys
import json
import argparse
from pathlib import Path
from itertools import cycle

import gen_data as G   # router _chat, provider cycle building blocks

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"
OUT = DATA / "route_train.jsonl"

# ── VERBATIM from internal/router/llmrouter.go (routeSystem). Keep in sync. ──
ROUTE_SYSTEM = """Классифицируй ровно одно сообщение пользователя. Верни ОДИН JSON-массив действий.

Ровно одно намерение: fact, reminder, note, query, act, chat, system.

Классифицируй по цели пользователя. Порядок решения:
1. Хочет напоминание в будущем → reminder
2. Явно просит сохранить информацию → note
3. Сообщает или обновляет текущее состояние/событие → fact
4. Хочет получить информацию → query
5. Просит выполнить работу → act
6. Про ассистента, настройки или память → system
7. Иначе → chat

Различия:
- note — сохранить информацию, без напоминания. text = суть.
- reminder — уведомить позже. text = что напомнить.
- fact — неявное обновление: пользователь сообщает, что что-то в мире изменилось (текущее/изменённое состояние, случившееся событие). key/value.

Примеры:
"запиши пароль" → {"intent":"note","text":"пароль"}
"напомни купить молоко" → {"intent":"reminder","text":"купить молоко"}
"запиши купить молоко" → {"intent":"note","text":"купить молоко"}
"я выпил воду" → {"intent":"fact","key":"water","value":"выпил"}
"мой любимый фильм — Интерстеллар" → {"intent":"note","text":"любимый фильм — Интерстеллар"}
"что такое docker?" → {"intent":"query","text":"что такое docker"}
"напиши письмо" → {"intent":"act","verb":"написать письмо"}
"очисти память" → {"intent":"system"}
"привет" → {"intent":"chat","text":"привет"}

Ответ — JSON-массив: по одному объекту на каждую просьбу. Обычно один. Если в реплике несколько просьб — по объекту на каждую. "напомни купить молоко, и запиши что кофе кончился" → [{"intent":"reminder","text":"купить молоко"},{"intent":"note","text":"кофе кончился"}]. Только JSON, без пояснений."""

INTENTS = {"fact", "reminder", "note", "query", "act", "chat", "system"}
ALLOWED_KEYS = {"intent", "key", "value", "text", "verb"}


def load_utterances() -> list[str]:
    """Real RU user turns to label. function_calling.jsonl is the main source;
    add user-*.jsonl fragments so chat/fact/system intents get coverage too."""
    seen, out = set(), []
    sources = ["function_calling.jsonl", "user-general.jsonl", "user-hello.jsonl",
               "user-who-are-you.jsonl"]
    for name in sources:
        p = DATA / name
        if not p.exists():
            continue
        for ln in p.read_text(encoding="utf-8").splitlines():
            ln = ln.strip()
            if not ln:
                continue
            try:
                obj = json.loads(ln)
                msgs = obj["messages"] if "messages" in obj else obj
                user = next((m["content"] for m in msgs if m.get("role") == "user"), None)
            except (json.JSONDecodeError, KeyError, TypeError):
                continue
            if user and user not in seen:
                seen.add(user)
                out.append(user)
    return out


def valid_action(obj) -> bool:
    return (isinstance(obj, dict)
            and obj.get("intent") in INTENTS
            and set(obj).issubset(ALLOWED_KEYS))


def valid_actions(arr) -> bool:
    """Contract is a JSON array of actions (compound utterance → N)."""
    return isinstance(arr, list) and len(arr) > 0 and all(valid_action(o) for o in arr)


def done_utterances() -> set:
    if not OUT.exists():
        return set()
    return {json.loads(l)["messages"][1]["content"]
            for l in OUT.read_text(encoding="utf-8").splitlines() if l.strip()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--check", action="store_true")
    args = ap.parse_args()

    if args.check:
        assert valid_action({"intent": "note", "text": "кофе закончился"})
        assert valid_action({"intent": "system"})
        assert not valid_action({"intent": "weather", "text": "x"})   # not in enum
        assert not valid_action({"intent": "note", "foo": "x"})       # stray key
        assert not valid_action({"text": "no intent"})
        assert valid_actions([{"intent": "reminder", "text": "x"}, {"intent": "note", "text": "y"}])
        assert not valid_actions([])                                   # empty array
        assert not valid_actions([{"intent": "note"}, {"intent": "bad"}])  # one bad
        assert not valid_actions({"intent": "note"})                  # not a list
        assert "JSON-массив" in ROUTE_SYSTEM
        assert "по одному объекту на каждую просьбу" in ROUTE_SYSTEM
        print("[✓] gen_route_data self-check passed")
        return

    utts = load_utterances()
    already = done_utterances()
    todo = [u for u in utts if u not in already]
    if args.limit:
        todo = todo[:args.limit]
    print(f"[*] utterances {len(utts)} | already {len(already)} | to do {len(todo)}")

    if args.dry_run:
        for u in todo[:10]:
            print(f"  {u}")
        print("[dry-run] no calls, no writes")
        return

    raw = os.environ.get("GEN_DATA_ROUTER_MODELS", "").strip()
    if not raw:
        sys.exit("set GEN_DATA_ROUTER_MODELS (comma-separated provider/model)")
    pairs = [s.strip().partition("/")[::2] for s in raw.split(",") if s.strip()]
    pcycle = cycle(pairs)
    router_url = os.environ.get("GEN_DATA_ROUTER_URL", "http://localhost:6446")

    written = 0
    with OUT.open("a", encoding="utf-8") as f:
        for i, utt in enumerate(todo):
            provider, model = next(pcycle)
            msgs = [{"role": "system", "content": ROUTE_SYSTEM},
                    {"role": "user", "content": utt}]
            content = G._chat(router_url, provider, model, msgs, max_tokens=128, temp=0.3)
            if content is None:
                continue
            try:
                arr = json.loads(content)
            except json.JSONDecodeError:
                continue
            if isinstance(arr, dict):        # tolerate a model that drops the wrapper
                arr = [arr]
            if not valid_actions(arr):
                continue
            f.write(json.dumps({
                "messages": [
                    {"role": "system", "content": ROUTE_SYSTEM},
                    {"role": "user", "content": utt},
                    {"role": "assistant", "content": json.dumps(arr, ensure_ascii=False)},
                ],
                "intents": [o["intent"] for o in arr],
                "source": "generated/route",
            }, ensure_ascii=False) + "\n")
            f.flush()
            written += 1
            if (i + 1) % 50 == 0:
                print(f"  {i+1}/{len(todo)} — {written} valid")

    print(f"[✓] wrote {written} route samples → {OUT}")


if __name__ == "__main__":
    main()