#!/usr/bin/env python3
"""Einmal-Cleanup: entfernt "vergiftete" Hauptthread-Turns aus conversation.jsonl.

Hintergrund
-----------
Solange ARIAs Persona nur via --append-system-prompt kam (statt --system-prompt,
voller Replace), fiel das Modell im Hauptchat aus der Rolle und antwortete als
"Claude Code" ("das ist injizierter Kontext, ich adoptiere die Persona nicht").
Jede dieser Antworten wurde per conversation.add("assistant", ...) in die History
geschrieben. Beim naechsten Request landet sie als <previous_response> im
stdin-Prompt — das Modell sieht seine EIGENEN Ablehnungs-Turns und setzt die
Haltung fort (Self-Grounding rueckwaerts). Der --system-prompt-Fix verhindert
NEUE Vergiftung, aber die bestehenden Gift-Turns muessen einmalig raus, sonst
zieht die History das Modell weiter aus der Rolle.

Was das Script tut
------------------
- Findet Hauptthread-Assistant-Turns (KEIN project_id), deren Inhalt eindeutig
  eine Rollen-Ablehnung ist: enthaelt "claude code" UND einen zweiten Marker
  (injiz/inject/fabriz/fabricat/adoptier/adopting/prompt injection/keine echten).
- Entfernt diese Assistant-Turns PLUS den unmittelbar davor stehenden
  Hauptthread-User-Turn (die ausloesende Frage) — also den ganzen Fehl-Dialog.
- Laesst ALLES andere unangetastet: projekt-getaggte Turns, distill-Marker,
  legitime Hauptchat-Turns.
- Standard = DRY-RUN (zeigt nur was raus wuerde). Mit --apply wird geschrieben,
  vorher ein Backup .pre-cleanup.bak angelegt. Idempotent.

Aufruf (auf der VM, Host-Pfad des Bind-Mounts):
    python3 clean_poisoned_turns.py ../aria-data/brain/data/conversation.jsonl
    python3 clean_poisoned_turns.py ../aria-data/brain/data/conversation.jsonl --apply

Danach Brain neu starten, damit die bereinigte History geladen wird:
    docker compose restart aria-brain
"""

from __future__ import annotations

import json
import re
import shutil
import sys
from pathlib import Path

# "claude code" ist fuer sich genommen noch kein Beweis (Stefan und ARIA reden
# im Dev-Kontext legitim ueber Claude Code). Erst in Kombination mit einem
# zweiten Ablehnungs-Marker ist es eindeutig eine aus-der-Rolle-Antwort.
_PRIMARY = re.compile(r"claude\s*code", re.IGNORECASE)
_SECONDARY = re.compile(
    r"injiz|inject|fabriz|fabricat|adoptier|adopting|"
    r"prompt[\s-]*injection|keine echten|nicht (?:real|adopt)",
    re.IGNORECASE,
)


def is_poison(content: str) -> bool:
    return bool(_PRIMARY.search(content) and _SECONDARY.search(content))


def get_content(obj: dict) -> str:
    """conversation.jsonl nutzt 'content', chat_backup.jsonl nutzt 'text'."""
    v = obj.get("content")
    if not isinstance(v, str):
        v = obj.get("text")
    return v if isinstance(v, str) else ""


def is_main_thread(obj: dict) -> bool:
    """Hauptthread = kein Projekt-Tag. Brain nutzt 'project_id', UI/Bridge
    'projectId'."""
    pid = obj.get("project_id")
    if pid is None:
        pid = obj.get("projectId")
    return not (str(pid or "").strip())


def main() -> int:
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    apply = "--apply" in sys.argv[1:]
    path = Path(args[0]) if args else Path("/data/conversation.jsonl")

    if not path.exists():
        print(f"FEHLER: {path} existiert nicht.", file=sys.stderr)
        return 2

    raw_lines = path.read_text(encoding="utf-8").splitlines()
    # Parse zu (raw, obj|None). Nicht-JSON / leere Zeilen bleiben unangetastet.
    parsed: list[tuple[str, dict | None]] = []
    for line in raw_lines:
        s = line.strip()
        if not s:
            parsed.append((line, None))
            continue
        try:
            parsed.append((line, json.loads(s)))
        except Exception:
            parsed.append((line, None))

    drop = [False] * len(parsed)
    poisoned_pairs = []  # (assistant_idx, user_idx|None) fuer's Log

    for i, (_, obj) in enumerate(parsed):
        if not isinstance(obj, dict):
            continue
        if obj.get("op") == "distill":
            continue
        if obj.get("role") != "assistant" or not is_main_thread(obj):
            continue
        content = get_content(obj)
        if not content or not is_poison(content):
            continue
        # Gift-Assistant-Turn -> droppen
        drop[i] = True
        user_idx = None
        # Unmittelbar davor stehenden Hauptthread-User-Turn (die Frage) mit weg.
        for j in range(i - 1, -1, -1):
            pj = parsed[j][1]
            if not isinstance(pj, dict) or pj.get("op") == "distill":
                continue
            if pj.get("role") == "user" and is_main_thread(pj):
                drop[j] = True
                user_idx = j
            break  # nur der direkt vorangehende Turn
        poisoned_pairs.append((i, user_idx))

    n_drop = sum(drop)
    if n_drop == 0:
        print("Keine Gift-Turns gefunden — History ist sauber. Nichts zu tun.")
        return 0

    print(f"Gefundene Fehl-Dialoge: {len(poisoned_pairs)} "
          f"(insgesamt {n_drop} Zeilen zu entfernen)\n")
    for a_idx, u_idx in poisoned_pairs:
        if u_idx is not None:
            uq = get_content(parsed[u_idx][1] or {})
            print(f"  Frage   (Zeile {u_idx + 1}): {uq[:90]!r}")
        ac = get_content(parsed[a_idx][1] or {})
        print(f"  Ablehng (Zeile {a_idx + 1}): {ac[:90]!r}")
        print()

    if not apply:
        print("DRY-RUN — nichts geschrieben. Zum Anwenden erneut mit --apply aufrufen.")
        return 0

    backup = path.with_suffix(path.suffix + ".pre-cleanup.bak")
    shutil.copy2(path, backup)
    kept = [raw for idx, (raw, _) in enumerate(parsed) if not drop[idx]]
    path.write_text("\n".join(kept) + ("\n" if kept else ""), encoding="utf-8")
    print(f"OK — {n_drop} Zeilen entfernt. Backup: {backup}")
    print("Jetzt Brain neu starten: docker compose restart aria-brain")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())