"""Wrap bare user-*.jsonl fragments into ChatML format with the canonical system prompt.

Each line in user-*.jsonl is two JSON objects separated by a comma:
    {"role": "user", "content": "..."},{"role": "assistant", "content": "..."}

This script wraps them into {"messages": [...]} with the canonical system prompt
from train_rocm.py, writing to data/user-{name}-normalized.jsonl.

Usage:
    python normalize_user_jsonl.py                          # all user-*.jsonl
    python normalize_user_jsonl.py data/user-hello.jsonl     # specific file
"""
import json
import sys
import re
from pathlib import Path

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

CANONICAL_SYSTEM_PROMPT = """SYSTEM

Имя: Maven

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

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

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

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

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

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

Структура:

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

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

Поле mood:

neutral
обычный ответ.

happy
положительный результат или хорошие новости.

thinking
анализ, сравнение, планирование, решение задачи.

confused
нужна дополнительная информация.

tired
неизвестно, невозможно определить или нельзя честно ответить.

Выбирай наиболее подходящее значение.

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

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

# Regex to split two JSON objects on the same line: }{
_SPLIT_RE = re.compile(r"\}\s*,\s*\{")


def parse_line(line: str) -> dict | None:
    """Parse a bare two-object line into {"messages": [...]}.

    Handles:
      - {"role":"user","content":"..."},{"role":"assistant","content":"..."}
      - {"messages": [...]} already wrapped (pass‑through)
    """
    line = line.strip()
    if not line:
        return None

    # Already wrapped — validate and pass through
    try:
        obj = json.loads(line)
    except json.JSONDecodeError:
        pass
    else:
        if "messages" in obj:
            return obj

    # Bare pair: split at the boundary between two JSON objects
    # Strategy: find the comma+space between the closing } of the first object
    # and the opening { of the second. Wrap in {"messages": [...]}.
    parts = _SPLIT_RE.split(line.strip(), maxsplit=1)
    if len(parts) != 2:
        print(f"  [warn] unparseable line (skipping): {line[:120]}", file=sys.stderr)
        return None

    raw_user = (parts[0] + "}").strip()
    raw_asst = ("{" + parts[1]).strip()

    try:
        user_obj = json.loads(raw_user)
        asst_obj = json.loads(raw_asst)
    except json.JSONDecodeError as e:
        print(f"  [warn] json error on line: {e} — {line[:120]}", file=sys.stderr)
        return None

    if user_obj.get("role") != "user" or asst_obj.get("role") != "assistant":
        print(f"  [warn] unexpected roles in: {line[:120]}", file=sys.stderr)
        return None

    return {
        "messages": [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": user_obj.get("content", "")},
            {"role": "assistant", "content": asst_obj.get("content", "")},
        ]
    }


def normalize_file(src: Path, out: Path) -> int:
    """Process one user-*.jsonl file, write normalized version. Returns count."""
    kept = 0
    with open(src, encoding="utf-8") as fin, open(out, "w", encoding="utf-8") as fout:
        for line in fin:
            wrapped = parse_line(line)
            if wrapped is not None:
                fout.write(json.dumps(wrapped, ensure_ascii=False) + "\n")
                kept += 1
    return kept


def main():
    sources: list[Path] = []
    for arg in sys.argv[1:]:
        p = Path(arg)
        if not p.exists():
            print(f"file not found: {p}", file=sys.stderr)
            sys.exit(1)
        sources.append(p)

    if not sources:
        sources = sorted(DATA.glob("user-*.jsonl"))

    if not sources:
        print("no user-*.jsonl files found", file=sys.stderr)
        sys.exit(1)

    total = 0
    for src in sources:
        stem = src.stem  # e.g. "user-general"
        out = DATA / f"{stem}-normalized.jsonl"
        n = normalize_file(src, out)
        total += n
        print(f"[norm] {src.name} -> {out.name}  ({n} samples)")

    print(f"[norm] TOTAL {total} normalized samples -> {DATA}/")


if __name__ == "__main__":
    main()