#!/usr/bin/env python3
"""
Filter exploration and self-play data for quality.

Keeps:
- Successful interactions (model looked up info AND executed correctly)
- First instance of each unique error pattern (for error correction training)
- High wiki-lookup-to-command ratios (model actually used the knowledge)

Removes:
- Duplicate topics (keeps first occurrence only)
- Empty responses (no commands, no lookups)
- Repeated failures on the same command pattern
- Rounds where model ignored wiki results

Output: data/processed/filtered_exploration.jsonl
"""

import json
import sys
from pathlib import Path
from collections import defaultdict

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent

EXPLORATION_DIR = PROJECT_ROOT / "data" / "raw" / "exploration_selfplay"
TOOL_SELFPLAY_DIR = PROJECT_ROOT / "data" / "raw" / "tool_selfplay"
OUTPUT_PATH = PROJECT_ROOT / "data" / "processed" / "filtered_exploration.jsonl"


def load_all_examples():
    """Load all exploration and tool self-play examples."""
    examples = []

    for jsonl in sorted(EXPLORATION_DIR.glob("*.jsonl")):
        with open(jsonl) as f:
            for line in f:
                if line.strip():
                    try:
                        examples.append(json.loads(line))
                    except json.JSONDecodeError:
                        pass

    for jsonl in sorted(TOOL_SELFPLAY_DIR.glob("*.jsonl")):
        with open(jsonl) as f:
            for line in f:
                if line.strip():
                    try:
                        examples.append(json.loads(line))
                    except json.JSONDecodeError:
                        pass

    return examples


def filter_examples(examples):
    """Filter for quality."""
    kept = []
    seen_topics = set()
    seen_error_patterns = set()
    stats = {
        "total": len(examples),
        "kept_success": 0,
        "kept_error_correction": 0,
        "kept_wiki_grounded": 0,
        "dropped_duplicate": 0,
        "dropped_empty": 0,
        "dropped_repeat_failure": 0,
    }

    for ex in examples:
        meta = ex.get("metadata", {})
        inp = ex.get("input", {})
        topic = inp.get("user_message", "")[:80]
        success_rate = meta.get("success_rate", meta.get("all_success", False))
        wiki_lookups = meta.get("wiki_lookups", 0)
        rcon_commands = meta.get("rcon_commands", 0)

        # Skip empty
        if rcon_commands == 0 and wiki_lookups == 0:
            stats["dropped_empty"] += 1
            continue

        # Deduplicate topics (keep first)
        if topic in seen_topics:
            stats["dropped_duplicate"] += 1
            continue
        seen_topics.add(topic)

        # Categorize
        if isinstance(success_rate, bool):
            is_success = success_rate
        else:
            is_success = success_rate > 0.7

        if is_success and rcon_commands > 0:
            # Successful interaction — always keep
            stats["kept_success"] += 1
            kept.append(ex)
        elif wiki_lookups > 0 and rcon_commands > 0:
            # Wiki-grounded (looked things up before acting) — keep even if some failures
            stats["kept_wiki_grounded"] += 1
            kept.append(ex)
        elif not is_success and rcon_commands > 0:
            # Failed — keep only first instance of each error pattern
            commands = ex.get("output", {}).get("commands", [])
            if commands:
                # Use first command as error pattern key
                pattern = commands[0][:40] if isinstance(commands[0], str) else ""
            else:
                rcon_results = meta.get("rcon_results", [])
                pattern = str(rcon_results[:1])[:60] if rcon_results else ""

            if pattern and pattern not in seen_error_patterns:
                seen_error_patterns.add(pattern)
                stats["kept_error_correction"] += 1
                kept.append(ex)
            else:
                stats["dropped_repeat_failure"] += 1
        else:
            stats["dropped_empty"] += 1

    return kept, stats


def main():
    print("Loading examples...")
    examples = load_all_examples()
    print(f"  Loaded {len(examples)} raw examples")

    print("Filtering...")
    filtered, stats = filter_examples(examples)

    print(f"\nFilter results:")
    for k, v in stats.items():
        print(f"  {k}: {v}")

    print(f"\nKept: {len(filtered)} ({100*len(filtered)//max(stats['total'],1)}%)")

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(OUTPUT_PATH, "w") as f:
        for ex in filtered:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    print(f"Written to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()
