#!/usr/bin/env python3
"""
filter_audit.py — Filter audit log data for 0.6.0 training.

Reads all audit sources (dev, prod, live, legacy audit_log), applies quality
filters, deduplicates, and writes a clean filtered_audit.jsonl.

Filters applied:
  - Drop god-mode with empty message (reinforces 30% empty-response bug)
  - Drop sudo-mode with empty commands
  - Drop prompt leaks (model echoed system prompt)
  - Drop enchantment syntax (Paper RCON rejects these)
  - Drop duplicates (keep first occurrence by user_message)
  - Convert legacy dict-format records to messages[] chat format

Usage:
    python3 training/scripts/filter_audit.py              # filter and write
    python3 training/scripts/filter_audit.py --dry-run    # show stats only
"""

import argparse
import json
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent.parent
RAW = ROOT / "data" / "raw"
PROCESSED = ROOT / "data" / "processed"
OUTPUT = PROCESSED / "filtered_audit.jsonl"

SOURCES = [
    RAW / "training_audit_dev_latest.jsonl",
    RAW / "training_audit_prod_latest.jsonl",
    RAW / "training_audit_live.jsonl",
    RAW / "audit_log.jsonl",
]

# Minimal system prompts for converting legacy dict-format records
SUDO_SYSTEM = (
    "/no_think\nYou are a Minecraft 1.21 command translator for a server admin. "
    "You receive natural language requests and return valid RCON commands.\n\n"
    "Return JSON: {\"commands\": [\"cmd1\", ...], \"reasoning\": \"why\"}"
)
GOD_SYSTEM = (
    "/no_think\nYou are God in a Minecraft server.\n\n"
    "Return JSON: {\"message\": \"Your dramatic response as God\", "
    "\"commands\": [\"cmd1\", \"cmd2\"], \"reasoning\": \"why\"}"
)


def is_prompt_leak(text: str) -> bool:
    """Detect if assistant response contains leaked system prompt fragments."""
    leak_signals = [
        "SYNTAX RULES",
        "command translator",
        "PERMISSION LEVEL",
        "CRITICAL RULES",
        "RISK GRADIENT",
        "Return ONLY JSON",
    ]
    hits = sum(1 for s in leak_signals if s in text)
    return hits >= 2


def has_enchantment_syntax(commands: list) -> bool:
    """Check if any command uses enchantment bracket syntax Paper rejects."""
    for cmd in commands:
        if re.search(r'\[enchantments=\{', str(cmd)):
            return True
        if re.search(r'\[potion_contents=\{', str(cmd)):
            return True
    return False


def convert_legacy(rec: dict) -> dict:
    """Convert old dict-format audit record to messages[] chat format."""
    mode = rec.get("mode", "sudo")
    player = rec.get("player", "unknown")
    query = rec.get("query", "")
    commands = rec.get("safe_commands") or rec.get("raw_commands", [])
    message = rec.get("message", "")
    reasoning = rec.get("reasoning", "")

    sys_prompt = GOD_SYSTEM if mode in ("god", "god_system") else SUDO_SYSTEM
    user_content = f"{'pray' if mode in ('god', 'god_system') else 'sudo'} {query}"

    response = {"commands": commands, "reasoning": reasoning}
    if mode in ("god", "god_system"):
        response["message"] = message or ""

    return {
        "messages": [
            {"role": "system", "content": sys_prompt},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": json.dumps(response)},
        ]
    }


def parse_assistant_response(msgs: list) -> tuple:
    """Extract (commands, message, raw_content) from assistant message."""
    for m in msgs:
        if m.get("role") != "assistant":
            continue
        content = m.get("content", "")
        try:
            resp = json.loads(content)
            return (
                resp.get("commands", []),
                resp.get("message", ""),
                content,
            )
        except (json.JSONDecodeError, TypeError):
            return ([], "", content)
    return ([], "", "")


def get_mode(msgs: list) -> str:
    """Determine if this is god or sudo mode from system prompt."""
    if not msgs:
        return "unknown"
    sys_content = msgs[0].get("content", "")
    if "You are God" in sys_content:
        return "god"
    if "command translator" in sys_content:
        return "sudo"
    return "unknown"


def get_user_message(msgs: list) -> str:
    """Extract user message content."""
    for m in msgs:
        if m.get("role") == "user":
            return m.get("content", "")
    return ""


def main():
    parser = argparse.ArgumentParser(description="Filter audit logs for 0.6.0 training")
    parser.add_argument("--dry-run", action="store_true", help="Show stats without writing")
    args = parser.parse_args()

    seen_user_msgs = set()
    kept = []
    dropped = {
        "god_empty_message": 0,
        "sudo_empty_commands": 0,
        "prompt_leak": 0,
        "enchantment_syntax": 0,
        "duplicate": 0,
        "parse_error": 0,
        "unknown_mode": 0,
    }
    source_counts = {}

    for source_path in SOURCES:
        if not source_path.exists():
            print(f"  Skip (not found): {source_path.name}")
            continue

        src_kept = 0
        src_total = 0

        with open(source_path) as f:
            for line in f:
                if not line.strip():
                    continue
                src_total += 1

                try:
                    rec = json.loads(line)
                except json.JSONDecodeError:
                    dropped["parse_error"] += 1
                    continue

                # Convert legacy dict format
                if isinstance(rec, dict) and "messages" not in rec:
                    rec = convert_legacy(rec)

                msgs = rec.get("messages", [])
                if not msgs:
                    dropped["parse_error"] += 1
                    continue

                mode = get_mode(msgs)
                user_msg = get_user_message(msgs)
                commands, message, raw_content = parse_assistant_response(msgs)

                # --- Filters ---

                if mode == "unknown":
                    dropped["unknown_mode"] += 1
                    continue

                # Deduplicate by user message
                msg_key = user_msg.strip()[:200]
                if msg_key in seen_user_msgs:
                    dropped["duplicate"] += 1
                    continue
                seen_user_msgs.add(msg_key)

                # Drop god-mode with empty message (reinforces the 30% empty bug)
                if mode == "god" and (not message or not message.strip()):
                    dropped["god_empty_message"] += 1
                    continue

                # Drop sudo-mode with empty commands
                if mode == "sudo" and not commands:
                    dropped["sudo_empty_commands"] += 1
                    continue

                # Drop prompt leaks
                if is_prompt_leak(raw_content):
                    dropped["prompt_leak"] += 1
                    continue

                # Drop enchantment bracket syntax
                if has_enchantment_syntax(commands):
                    dropped["enchantment_syntax"] += 1
                    continue

                kept.append(rec)
                src_kept += 1

        source_counts[source_path.name] = (src_total, src_kept)
        print(f"  {source_path.name}: {src_total} total → {src_kept} kept")

    # Summary
    total_input = sum(v[0] for v in source_counts.values())
    print(f"\n{'='*60}")
    print(f"Total input:  {total_input}")
    print(f"Total kept:   {len(kept)}")
    print(f"Total dropped: {total_input - len(kept)}")
    print(f"\nDrop reasons:")
    for reason, count in sorted(dropped.items(), key=lambda x: -x[1]):
        if count > 0:
            print(f"  {reason:25} {count:6}")

    # Mode breakdown of kept
    god_kept = sum(1 for r in kept if get_mode(r.get("messages", [])) == "god")
    sudo_kept = sum(1 for r in kept if get_mode(r.get("messages", [])) == "sudo")
    print(f"\nKept breakdown: {god_kept} god + {sudo_kept} sudo")

    if args.dry_run:
        print(f"\n[DRY RUN] Would write {len(kept)} records to {OUTPUT}")
        return

    # Write output
    with open(OUTPUT, "w") as f:
        for rec in kept:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")

    print(f"\nWrote {len(kept)} records to {OUTPUT}")


if __name__ == "__main__":
    main()
