#!/usr/bin/env python3
"""
Filter audit log data for training quality.

Keeps the full validator loop (generated → executed → RCON result) as the
training signal. Quarantines empty outputs, system prompt leaks, and broken
JSON/tellraw commands into a separate file for review.

Usage:
    python3 filter_audit_log.py [--input FILE] [--output FILE] [--stats]
    python3 filter_audit_log.py --stats  # dry run, print breakdown only
"""

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
DEFAULT_INPUT = PROJECT_ROOT / "data" / "raw" / "training_audit_dev_latest.jsonl"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "processed" / "filtered_audit.jsonl"
DEFAULT_QUARANTINE = PROJECT_ROOT / "data" / "quarantine" / "audit_quarantine.jsonl"

# Strings that indicate system prompt leakage in model output
PROMPT_LEAK_PATTERNS = [
    "you are a minecraft",
    "command translator",
    "player request vector",
    "you are an ai",
    "system prompt",
    "Context: Player is",
]


def is_prompt_leak(message: str) -> bool:
    msg_lower = message.lower()
    return any(pat in msg_lower for pat in PROMPT_LEAK_PATTERNS)


def has_broken_json(commands: list[str]) -> bool:
    """Detect broken tellraw/JSON commands (bad escaping, fragments)."""
    for cmd in commands:
        stripped = cmd.strip()
        # Fragments from split JSON: starts with colon, comma, or brace-colon
        if stripped.startswith((":\\", ",\\", ":{", ",")):
            return True
        # Truncated tellraw with unmatched braces
        if "tellraw" in stripped:
            opens = stripped.count("{") + stripped.count("[")
            closes = stripped.count("}") + stripped.count("]")
            if abs(opens - closes) > 1:
                return True
    return False


def classify(record: dict) -> tuple[str, str]:
    """Classify a record as (keep|quarantine, reason)."""
    mode = record.get("mode", "")
    output = record.get("output", {})
    message = output.get("message", "").strip()
    cmds_gen = output.get("commands_generated", [])
    cmds_exe = output.get("commands_executed", [])

    # ── Quarantine ─────────────────────────────────────────────────
    # System prompt leak
    if is_prompt_leak(message):
        return "quarantine", "prompt_leak"

    # Completely empty output — teaches the model to produce nothing
    if not message and not cmds_gen:
        return "quarantine", "empty_output"

    # Broken JSON/tellraw commands — teaches bad syntax
    if cmds_gen and has_broken_json(cmds_gen):
        return "quarantine", "broken_json"

    # ── Keep ───────────────────────────────────────────────────────
    # Validator correction: generated != executed (highest value)
    if cmds_gen and cmds_exe and cmds_gen != cmds_exe:
        return "keep", "validator_corrected"

    # Generated but validator blocked — teaches what gets rejected
    if cmds_gen and not cmds_exe:
        return "keep", "validator_blocked"

    # Commands executed (sudo or god mode with action)
    if cmds_gen and cmds_exe:
        return "keep", "executed"

    # God mode with RP message, no commands — valid roleplay response
    if mode == "god" and message:
        return "keep", "god_rp"

    # Catch-all: keep with flag
    return "keep", "other"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--quarantine", type=Path, default=DEFAULT_QUARANTINE,
                        help="Write quarantined records here for review")
    parser.add_argument("--stats", action="store_true", help="Print stats only, don't write")
    args = parser.parse_args()

    stats = {"keep": {}, "quarantine": {}}
    kept = []
    quarantined = []

    with open(args.input) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            action, reason = classify(record)
            stats[action][reason] = stats[action].get(reason, 0) + 1
            record["_filter_action"] = action
            record["_filter_reason"] = reason
            if action == "keep":
                kept.append(record)
            else:
                quarantined.append(record)

    total = len(kept) + len(quarantined)
    print(f"Total records:    {total}")
    print(f"Kept:             {len(kept)} ({100*len(kept)/total:.1f}%)")
    print(f"Quarantined:      {len(quarantined)} ({100*len(quarantined)/total:.1f}%)")
    print()

    print("KEPT breakdown:")
    for reason, count in sorted(stats["keep"].items(), key=lambda x: -x[1]):
        print(f"  {reason:<25} {count:>5}")
    print()
    print("QUARANTINED breakdown:")
    for reason, count in sorted(stats["quarantine"].items(), key=lambda x: -x[1]):
        print(f"  {reason:<25} {count:>5}")

    if args.stats:
        return

    # Write kept records (strip internal filter tags)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w") as f:
        for record in kept:
            out = {k: v for k, v in record.items() if not k.startswith("_filter_")}
            f.write(json.dumps(out) + "\n")
    print(f"\nWrote {len(kept)} records to {args.output}")

    # Write quarantined records (keep filter tags for review)
    args.quarantine.parent.mkdir(parents=True, exist_ok=True)
    with open(args.quarantine, "w") as f:
        for record in quarantined:
            f.write(json.dumps(record) + "\n")
    print(f"Wrote {len(quarantined)} quarantined records to {args.quarantine}")


if __name__ == "__main__":
    main()
