"""Phase 6b — generate failure samples (tired/confused) and append to persona_train.jsonl.

Borrows CANONICAL_SYSTEM_PROMPT, failure prompts, and validation from gen_data.py.
Runs independently so we can iterate on failure prompting without touching the
persona generation pipeline.

Usage:
    python append_failures.py

Environment:
    INFERENCE_ROUTER_URL  (default: http://inference.kvmx.ru)
"""
import json
import os
import sys
import random
import re
from itertools import cycle
from pathlib import Path

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

random.seed(42)

# steal the canonical system prompt and failure prompts from gen_data.py
sys.path.insert(0, str(HERE))
import gen_data  # noqa: E402

CANONICAL_SYSTEM_PROMPT = gen_data.CANONICAL_SYSTEM_PROMPT
FAILURE_PROMPTS_TIRED = gen_data.FAILURE_PROMPTS_TIRED
FAILURE_PROMPTS_CONFUSED = gen_data.FAILURE_PROMPTS_CONFUSED
VALID_MOODS = gen_data.VALID_MOODS
validate_sample = gen_data.validate_sample

# ── router config ────────────────────────────────────────────────────────
ROUTER_URL = os.environ.get(
    "INFERENCE_ROUTER_URL", "http://inference.kvmx.ru"
).rstrip("/")

# providers to cycle through — brief list of models that work well for RU
ROUTER_PAIRS = [
    # format: (provider, model) — X-Provider header + bare model name
    ("cerebras", "gemma-4-31b"),
    ("cerebras", "zai-glm-4.7"),
    ("github_models", "models/meta/llama-4"),
    ("github_models", "models/mistral/mistral-large-3"),
    ("mistral", "mistral-large-latest"),
    ("mistral", "codestral-latest"),
]


# ── router call ──────────────────────────────────────────────────────────
def _chat(provider: str, model: str, msgs: list,
          max_tokens: int = 200, temp: float = 0.7) -> str | None:
    import time as _time
    import urllib.request as _req
    import urllib.error as _err

    base = ROUTER_URL
    if not base.endswith("/v1"):
        base += "/v1"
    url = f"{base}/chat/completions"
    headers = {"Content-Type": "application/json"}
    headers["X-Provider"] = provider

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

    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()
            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:
            _time.sleep(min(2 ** attempt, 8))
        except Exception:
            _time.sleep(1)
    return None


def _parse_asst(content: str) -> dict | 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


def flatten_msgs(user_text: str, asst_text: str) -> list[dict]:
    return [
        {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
        {"role": "user", "content": user_text},
        {"role": "assistant", "content": asst_text},
    ]


# ── generation ───────────────────────────────────────────────────────────
def gen_failures(count: int) -> list[dict]:
    """Generate tired + confused samples, return list of validated dicts."""
    samples = []
    providers = cycle(ROUTER_PAIRS)

    # tired: 2/3 of count
    tired_target = count * 2 // 3
    pool_tired = FAILURE_PROMPTS_TIRED * (tired_target // len(FAILURE_PROMPTS_TIRED) + 1)
    random.shuffle(pool_tired)

    print(f"  tired: targeting {tired_target} samples...")
    done = 0
    for prompt_text in pool_tired[:tired_target]:
        provider, model = next(providers)
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": (
                f"{prompt_text}\n\n"
                f"(Это вопрос, на который у тебя нет ответа. "
                f"Не выдумывай. Честно скажи, что не знаешь или не можешь "
                f"ответить. Используй mood: tired. Только tired, не neutral, "
                f"не confused, не thinking — именно tired.)"
            )},
        ]
        content = _chat(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",
        })
        done += 1
        if done % 25 == 0:
            print(f"    ...{done}/{tired_target} tired valid")

    print(f"  tired: {done} valid")

    # confused: 1/3 of count
    confused_target = count // 3
    pool_confused = FAILURE_PROMPTS_CONFUSED * (confused_target // len(FAILURE_PROMPTS_CONFUSED) + 1)
    random.shuffle(pool_confused)

    print(f"  confused: targeting {confused_target} samples...")
    done = 0
    for prompt_text in pool_confused[:confused_target]:
        provider, model = next(providers)
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": (
                f"{prompt_text}\n\n"
                f"(Это неясный или слишком общий запрос. Ответь коротким "
                f"уточняющим вопросом. Используй mood: confused. Только "
                f"confused, не neutral, не tired, не thinking — именно confused.)"
            )},
        ]
        content = _chat(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",
        })
        done += 1
        if done % 25 == 0:
            print(f"    ...{done}/{confused_target} confused valid")

    print(f"  confused: {done} valid")
    return samples


def main():
    # target: 750 failure samples (15% of 5000)
    target = 750
    print(f"Generating up to {target} failure samples via router...")
    print(f"  Router: {ROUTER_URL}")
    print(f"  Models: {[f'{p}/{m}' for p,m in ROUTER_PAIRS]}")

    samples = gen_failures(target)
    print(f"\nGot {len(samples)} valid failure samples")

    if not samples:
        print("No samples generated — nothing to append.")
        return

    # dedup by response text against existing train data
    train_path = DATA / "persona_train.jsonl"
    if train_path.exists():
        existing_responses: set[str] = set()
        with open(train_path) as f:
            for line in f:
                try:
                    d = json.loads(line)
                    msg = d.get("messages", [])
                    if len(msg) >= 3:
                        asst_raw = json.loads(msg[2]["content"])
                        existing_responses.add(asst_raw.get("response", "").strip().lower())
                except (json.JSONDecodeError, KeyError, IndexError):
                    pass
        before = len(samples)
        samples = [s for s in samples if s["response"].strip().lower() not in existing_responses]
        print(f"Dedup removed {before - len(samples)} (already in train file)")

    # dedup within new samples
    seen: set[str] = set()
    deduped = []
    for s in samples:
        key = s["response"].strip().lower()
        if key not in seen:
            seen.add(key)
            deduped.append(s)
    samples = deduped

    # append to persona_train.jsonl
    with open(train_path, "a", encoding="utf-8") as f:
        for s in samples:
            f.write(json.dumps({"messages": s["messages"]}, ensure_ascii=False) + "\n")

    print(f"Appended {len(samples)} failure samples to {train_path}")
    print(f"Mood distribution: tired={sum(1 for s in samples if s['mood']=='tired')}, "
          f"confused={sum(1 for s in samples if s['mood']=='confused')}")


if __name__ == "__main__":
    main()