#!/usr/bin/env python3
"""
Generate high-quality tool-calling training data with REAL RCON responses.

No AI involved — we craft the prompts, commands, and tool sequences by hand,
then execute each command through RCON to get actual server responses.
This produces gold-standard training data with real validation.

Generates examples across all 14 tools with proper multi-turn conversations:
  - script.validate → script.write → script.execute (with real RCON validation)
  - memory.read → rcon.execute (tp to saved location)
  - world.nearby_entities → rcon.execute (kill scanned mobs)
  - wiki_lookup → rcon.execute (apply looked-up knowledge)
  - chained multi-tool sequences

Usage:
    python3 generate_rcon_validated_training.py --rcon-host 192.168.0.244 --rcon-port 25577
"""

import argparse
import json
import random
import re
import socket
import struct
import sys
import time
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from agent.tools.tool_schemas import qwen3_tools_block
from agent.prompts.system_prompts import SYNTAX_RULES, RISK_GRADIENT

OUTPUT_PATH = PROJECT_ROOT / "data" / "raw" / "rcon_validated_tool_training.jsonl"

TOOLS_BLOCK = qwen3_tools_block()
SYSTEM_SUDO = (
    "You are a Minecraft 1.21 command translator with full tool access.\n"
    "Tools: rcon.execute, minecraft.wiki_lookup, world.player_info, world.server_state, "
    "world.nearby_entities, memory.read, memory.write, script.write, script.validate, "
    "script.execute, script.read, script.list, script.delete, script.schedule.\n\n"
    "For complex builds (4+ commands), write mcfunction scripts. Validate before writing.\n"
    "PERMISSION LEVEL: 4 (generous).\n\n"
    "Return JSON: {\"risk_level\": <0-5>, \"commands\": [...], \"reasoning\": \"...\"}\n\n"
    + SYNTAX_RULES + RISK_GRADIENT + "\n" + TOOLS_BLOCK
)

SYSTEM_GOD = (
    "You are God in a Minecraft server. You have full tool access.\n"
    "Return JSON: {\"risk_level\": <0-5>, \"message\": \"...\", \"commands\": [...], \"reasoning\": \"...\"}\n\n"
    + SYNTAX_RULES + "\n" + TOOLS_BLOCK
)

PLAYERS = ["slingshooter08", "Ace13245", "TheBigBoss", "xXDragonSlayerXx"]


from agent.tools.persistent_rcon import get_rcon as _get_rcon

class SimpleRCON:
    """Wrapper around persistent_rcon."""
    def __init__(self, host, port, password):
        self._rcon = _get_rcon(host, port, password)

    def command(self, cmd):
        return self._rcon.command(cmd)


def sys_msg(mode="sudo"):
    return {"role": "system", "content": SYSTEM_GOD if mode == "god" else SYSTEM_SUDO}

def user_msg(text):
    return {"role": "user", "content": text}

def tool_call(name, args):
    return {"role": "assistant", "content": f"<tool_call>\n{json.dumps({'name': name, 'arguments': args})}\n</tool_call>"}

def tool_result(data):
    return {"role": "tool", "content": json.dumps(data)}

def final_response(resp):
    return {"role": "assistant", "content": json.dumps(resp)}

def is_error(result):
    return any(e in result for e in ("<--[HERE]", "Unknown", "Incorrect", "Expected", "Invalid"))


def validate_commands_rcon(rcon, commands):
    """Execute commands through RCON and return results."""
    results = []
    for cmd in commands:
        cmd = cmd.strip()
        if not cmd or cmd.startswith("#"):
            results.append({"cmd": cmd, "result": "comment", "ok": True})
            continue
        try:
            result = rcon.command(cmd)
            ok = not is_error(result)
            results.append({"cmd": cmd, "result": result[:300], "ok": ok})
        except Exception as e:
            results.append({"cmd": cmd, "result": str(e), "ok": False})
    return results


def gen_script_examples(rcon):
    """Script write/validate/execute examples with real RCON validation."""
    examples = []

    SCRIPTS = [
        ("build me a small oak cabin", "oak_cabin", "Small oak cabin with door, windows, and lantern", [
            "fill ~-3 ~ ~-3 ~3 ~3 ~3 minecraft:oak_planks hollow",
            "fill ~-3 ~ ~-3 ~3 ~ ~3 minecraft:oak_planks",
            "fill ~-2 ~1 ~-3 ~2 ~2 ~-3 minecraft:air",
            "setblock ~0 ~1 ~-3 minecraft:oak_door[facing=north,half=lower]",
            "setblock ~0 ~2 ~-3 minecraft:oak_door[facing=north,half=upper]",
            "setblock ~2 ~1 ~3 minecraft:glass_pane",
            "setblock ~-2 ~1 ~3 minecraft:glass_pane",
            "setblock ~0 ~2 ~0 minecraft:lantern[hanging=true]",
        ]),
        ("create a pvp arena with wool corners", "pvp_arena", "PvP arena with colored corners", [
            "fill ~-12 ~-1 ~-12 ~12 ~-1 ~12 minecraft:smooth_stone",
            "fill ~-12 ~ ~-12 ~12 ~4 ~-12 minecraft:iron_bars",
            "fill ~-12 ~ ~12 ~12 ~4 ~12 minecraft:iron_bars",
            "fill ~-12 ~ ~-12 ~-12 ~4 ~12 minecraft:iron_bars",
            "fill ~12 ~ ~-12 ~12 ~4 ~12 minecraft:iron_bars",
            "fill ~-12 ~-1 ~-12 ~-10 ~-1 ~-10 minecraft:red_wool",
            "fill ~10 ~-1 ~10 ~12 ~-1 ~12 minecraft:blue_wool",
        ]),
        ("make an enchanting room", "enchant_room", "Enchanting table surrounded by bookshelves", [
            "fill ~-3 ~ ~-3 ~3 ~3 ~3 minecraft:deepslate_bricks hollow",
            "fill ~-2 ~ ~-2 ~2 ~2 ~2 minecraft:air",
            "setblock ~0 ~ ~0 minecraft:enchanting_table",
            "fill ~-2 ~1 ~-2 ~2 ~1 ~2 minecraft:bookshelf",
            "fill ~-1 ~1 ~-1 ~1 ~1 ~1 minecraft:air",
            "setblock ~0 ~1 ~0 minecraft:enchanting_table",
            "setblock ~0 ~2 ~0 minecraft:lantern[hanging=true]",
        ]),
        ("build a nether portal frame", "nether_portal", "Obsidian nether portal frame", [
            "fill ~0 ~ ~0 ~3 ~4 ~0 minecraft:obsidian",
            "fill ~1 ~1 ~0 ~2 ~3 ~0 minecraft:air",
        ]),
        ("create a mob farm collection area", "mob_farm", "Mob farm with hoppers and chest", [
            "fill ~-4 ~-1 ~-4 ~4 ~-1 ~4 minecraft:cobblestone",
            "setblock ~0 ~-2 ~0 minecraft:chest",
            "setblock ~0 ~-1 ~0 minecraft:hopper",
            "fill ~-4 ~ ~0 ~-1 ~ ~0 minecraft:water",
            "fill ~1 ~ ~0 ~4 ~ ~0 minecraft:water",
        ]),
        ("build a lookout tower", "lookout_tower", "Tall stone tower with viewing platform", [
            "fill ~-1 ~ ~-1 ~1 ~8 ~1 minecraft:stone_bricks hollow",
            "fill ~0 ~ ~0 ~0 ~7 ~0 minecraft:air",
            "fill ~1 ~1 ~0 ~1 ~7 ~0 minecraft:ladder[facing=west]",
            "fill ~-2 ~8 ~-2 ~2 ~8 ~2 minecraft:stone_brick_slab",
            "fill ~-2 ~9 ~-2 ~2 ~9 ~2 minecraft:stone_brick_wall",
            "fill ~-1 ~9 ~-1 ~1 ~9 ~1 minecraft:air",
            "setblock ~0 ~10 ~0 minecraft:lantern",
        ]),
        ("make a fishing dock", "fishing_dock", "Wooden dock extending over water", [
            "fill ~0 ~-1 ~0 ~1 ~-1 ~8 minecraft:oak_planks",
            "fill ~-1 ~-2 ~0 ~-1 ~-1 ~8 minecraft:oak_fence",
            "fill ~2 ~-2 ~0 ~2 ~-1 ~8 minecraft:oak_fence",
            "setblock ~0 ~0 ~8 minecraft:oak_fence",
            "setblock ~1 ~0 ~8 minecraft:oak_fence",
            "setblock ~0 ~0 ~0 minecraft:lantern",
        ]),
        ("create a garden with flower beds", "garden", "Flower garden with paths", [
            "fill ~-5 ~-1 ~-5 ~5 ~-1 ~5 minecraft:grass_block",
            "fill ~-5 ~-1 ~0 ~5 ~-1 ~0 minecraft:gravel",
            "fill ~0 ~-1 ~-5 ~0 ~-1 ~5 minecraft:gravel",
            "setblock ~-3 ~ ~-3 minecraft:rose_bush[half=lower]",
            "setblock ~3 ~ ~-3 minecraft:peony[half=lower]",
            "setblock ~-3 ~ ~3 minecraft:lilac[half=lower]",
            "setblock ~3 ~ ~3 minecraft:sunflower[half=lower]",
            "setblock ~0 ~ ~0 minecraft:water",
        ]),
        ("build a redstone clock", "redstone_clock", "Simple repeater clock", [
            "setblock ~0 ~ ~0 minecraft:redstone_wire",
            "setblock ~1 ~ ~0 minecraft:repeater[facing=west,delay=4]",
            "setblock ~2 ~ ~0 minecraft:redstone_wire",
            "setblock ~2 ~ ~1 minecraft:repeater[facing=north,delay=4]",
            "setblock ~2 ~ ~2 minecraft:redstone_wire",
            "setblock ~1 ~ ~2 minecraft:repeater[facing=east,delay=4]",
            "setblock ~0 ~ ~2 minecraft:redstone_wire",
            "setblock ~0 ~ ~1 minecraft:repeater[facing=south,delay=4]",
        ]),
        ("make a treasure vault", "treasure_vault", "Secure room with chests and iron door", [
            "fill ~-3 ~ ~-3 ~3 ~3 ~3 minecraft:deepslate_brick_wall hollow",
            "fill ~-2 ~ ~-2 ~2 ~2 ~2 minecraft:air",
            "setblock ~0 ~0 ~-3 minecraft:iron_door[facing=north,half=lower]",
            "setblock ~0 ~1 ~-3 minecraft:iron_door[facing=north,half=upper]",
            "setblock ~-1 ~0 ~-2 minecraft:chest",
            "setblock ~1 ~0 ~-2 minecraft:chest",
            "setblock ~-1 ~0 ~2 minecraft:chest",
            "setblock ~1 ~0 ~2 minecraft:chest",
            "setblock ~0 ~2 ~0 minecraft:lantern[hanging=true]",
        ]),
    ]

    for prompt, name, desc, commands in SCRIPTS:
        player = random.choice(PLAYERS)
        mode = "sudo"
        msgs = [sys_msg(mode), user_msg(f"Player {player}: sudo {prompt}")]

        # Step 1: Validate
        rcon_results = validate_commands_rcon(rcon, commands)
        valid = all(r["ok"] for r in rcon_results)
        errors = [{"line": i+1, "command": r["cmd"], "error": r["result"]}
                  for i, r in enumerate(rcon_results) if not r["ok"]]

        msgs.append(tool_call("script.validate", {"commands": commands}))
        msgs.append(tool_result({
            "valid": valid,
            "total": len(commands),
            "passed": sum(1 for r in rcon_results if r["ok"]),
            "errors": errors,
        }))

        # Step 2: Write (if valid or partially valid)
        write_cmds = [r["cmd"] for r in rcon_results if r["ok"]] if not valid else commands
        msgs.append(tool_call("script.write", {"name": name, "commands": write_cmds, "description": desc}))
        msgs.append(tool_result({"ok": True, "path": f"mortdecai:{name}", "lines": len(write_cmds)}))

        # Step 3: Execute
        msgs.append(tool_call("script.execute", {"name": name, "as_player": player}))
        msgs.append(tool_result({"ok": True, "result": f"Executed {len(write_cmds)} commands from function mortdecai:{name}"}))

        # Final
        fix_note = f" Fixed {len(errors)} validation errors." if errors else ""
        resp = {
            "risk_level": 3,
            "commands": [f"function mortdecai:{name}"],
            "reasoning": f"Validated {len(commands)} commands ({sum(1 for r in rcon_results if r['ok'])} passed).{fix_note} Wrote and executed '{name}' script at player position.",
        }
        msgs.append(final_response(resp))

        examples.append({
            "id": f"rcon-script-{name}",
            "source": "rcon_validated",
            "type": "script_write_execute",
            "messages": msgs,
        })

    return examples


def gen_direct_command_examples(rcon):
    """Direct rcon.execute examples with real server responses."""
    examples = []

    COMMANDS = [
        ("sudo give me a diamond sword with sharpness 5", [
            "give {p} minecraft:diamond_sword[enchantments={sharpness:5,unbreaking:3,mending:1}] 1",
        ], "Gave max enchanted diamond sword."),
        ("sudo full netherite armor please", [
            "give {p} minecraft:netherite_helmet[enchantments={protection:4,unbreaking:3,mending:1}] 1",
            "give {p} minecraft:netherite_chestplate[enchantments={protection:4,unbreaking:3,mending:1}] 1",
            "give {p} minecraft:netherite_leggings[enchantments={protection:4,unbreaking:3,mending:1}] 1",
            "give {p} minecraft:netherite_boots[enchantments={protection:4,unbreaking:3,mending:1}] 1",
        ], "Full netherite armor set with protection 4 and mending."),
        ("sudo make it night", ["time set night"], "Set time to night."),
        ("sudo give me 64 golden apples", [
            "give {p} minecraft:golden_apple 64",
        ], "Gave 64 golden apples."),
        ("sudo smite me with lightning", [
            "execute at {p} run summon minecraft:lightning_bolt",
        ], "Summoned lightning at player position."),
        ("sudo give me speed 3 for 10 minutes", [
            "effect give {p} minecraft:speed 600 2",
        ], "Applied speed 3 (amplifier 2) for 600 seconds."),
        ("sudo clear weather", ["weather clear 999999"], "Cleared weather."),
        ("sudo make me invincible", [
            "effect give {p} minecraft:resistance 6000 4",
            "effect give {p} minecraft:regeneration 6000 2",
            "effect give {p} minecraft:fire_resistance 6000 0",
        ], "Applied resistance 5, regeneration 3, and fire resistance for 100 minutes."),
        ("sudo give me a god bow", [
            "give {p} minecraft:bow[enchantments={power:5,infinity:1,flame:1,punch:2,unbreaking:3}] 1",
        ], "Bow with power 5, infinity, flame, punch 2."),
        ("sudo stack of each wood type", [
            "give {p} minecraft:oak_log 64",
            "give {p} minecraft:spruce_log 64",
            "give {p} minecraft:birch_log 64",
            "give {p} minecraft:jungle_log 64",
            "give {p} minecraft:acacia_log 64",
            "give {p} minecraft:dark_oak_log 64",
            "give {p} minecraft:cherry_log 64",
            "give {p} minecraft:mangrove_log 64",
        ], "All 8 wood types, one stack each."),
        ("sudo clear my inventory", ["clear {p}"], "Cleared all items."),
        ("sudo tp me to 0 100 0", ["tp {p} 0 100 0"], "Teleported to coordinates."),
        ("sudo gamemode creative", ["gamemode creative {p}"], "Set creative mode."),
        ("sudo gamemode survival", ["gamemode survival {p}"], "Set survival mode."),
        ("sudo kill all zombies", ["kill @e[type=minecraft:zombie]"], "Killed all zombies."),
        ("sudo summon 5 cows", [
            "summon minecraft:cow ~ ~ ~1",
            "summon minecraft:cow ~ ~ ~2",
            "summon minecraft:cow ~ ~ ~3",
            "summon minecraft:cow ~1 ~ ~1",
            "summon minecraft:cow ~1 ~ ~2",
        ], "Summoned 5 cows nearby."),
        ("sudo give everyone a diamond", [
            "give @a minecraft:diamond 1",
        ], "Gave all players a diamond."),
        ("sudo set time to sunrise", ["time set 23000"], "Set to dawn (tick 23000)."),
        ("sudo fill a 5x5 area with gold blocks under me", [
            "fill ~-2 ~-1 ~-2 ~2 ~-1 ~2 minecraft:gold_block",
        ], "Filled 5x5 gold platform under player."),
        ("sudo give me a trident with loyalty and channeling", [
            "give {p} minecraft:trident[enchantments={loyalty:3,channeling:1,impaling:5,unbreaking:3,mending:1}] 1",
        ], "Trident with loyalty 3, channeling, impaling 5. Loyalty and riptide are mutually exclusive — used loyalty."),
        ("sudo night vision forever", [
            "effect give {p} minecraft:night_vision 999999 0",
        ], "Permanent night vision."),
        ("sudo repair what I'm holding", [
            "give {p} minecraft:experience_bottle 64",
        ], "Gave XP bottles — use with mending items to repair. Cannot directly repair via commands."),
        ("sudo give me a mace", [
            "give {p} minecraft:mace[enchantments={density:5,wind_burst:3,unbreaking:3,mending:1}] 1",
        ], "Mace with density 5 and wind burst 3. New 1.21 weapon."),
        ("sudo spawn a warden", [
            "summon minecraft:warden ~ ~ ~5",
        ], "Summoned warden 5 blocks away. Be careful — 500 HP, 2-hit kill."),
        ("sudo give me all the music discs", [
            "give {p} minecraft:music_disc_13 1",
            "give {p} minecraft:music_disc_cat 1",
            "give {p} minecraft:music_disc_blocks 1",
            "give {p} minecraft:music_disc_chirp 1",
            "give {p} minecraft:music_disc_far 1",
            "give {p} minecraft:music_disc_mall 1",
            "give {p} minecraft:music_disc_mellohi 1",
            "give {p} minecraft:music_disc_stal 1",
            "give {p} minecraft:music_disc_strad 1",
            "give {p} minecraft:music_disc_ward 1",
            "give {p} minecraft:music_disc_pigstep 1",
            "give {p} minecraft:music_disc_otherside 1",
        ], "All music discs."),
        ("sudo potion of the turtle master", [
            "give {p} minecraft:potion[potion_contents={potion:\"minecraft:turtle_master\"}] 1",
        ], "Turtle master potion — slowness 4 + resistance 3."),
        ("sudo banner with a skull pattern", [
            "give {p} minecraft:white_banner[banner_patterns=[{pattern:\"minecraft:skull\",color:\"black\"}]] 1",
        ], "White banner with black skull pattern using 1.21 syntax."),
    ]

    for prompt, cmds_tmpl, reasoning in COMMANDS:
        player = random.choice(PLAYERS)
        cmds = [c.replace("{p}", player) for c in cmds_tmpl]
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {prompt}")]

        all_ok = True
        for cmd in cmds:
            rcon_result = rcon.command(cmd)
            ok = not is_error(rcon_result)
            if not ok:
                all_ok = False
            msgs.append(tool_call("rcon.execute", {"command": cmd}))
            msgs.append(tool_result({"success": ok, "result": rcon_result[:300]}))

        resp = {"risk_level": 3, "commands": cmds, "reasoning": reasoning}
        msgs.append(final_response(resp))

        examples.append({
            "id": f"rcon-direct-{len(examples):03d}",
            "source": "rcon_validated",
            "type": "direct_command",
            "rcon_validated": True,
            "all_success": all_ok,
            "messages": msgs,
        })

    return examples


def gen_memory_examples(rcon):
    """Memory read/write examples."""
    examples = []

    MEMORY_OPS = [
        ("sudo remember this as home", "write", "home", "location", {"x": 100, "y": 64, "z": -200}, None,
         "Saved current location as 'home'."),
        ("sudo save this spot as my base", "write", "base", "location", {"x": -500, "y": 72, "z": 300}, None,
         "Saved base location."),
        ("sudo tp me home", "read", "home", None, None, "tp {p} 100 64 -200",
         "Read 'home' from memory. Teleporting."),
        ("sudo take me to my base", "read", "base", None, None, "tp {p} -500 72 300",
         "Read 'base' from memory. Teleporting."),
        ("sudo what do you remember about me?", "read", None, None, None, None,
         "Listed all memories for this player."),
        ("sudo remember my favorite tool is a pickaxe", "write", "favorite_tool", "preference", "diamond pickaxe", None,
         "Saved tool preference."),
        ("sudo remember Ace is my enemy", "write", "enemy_ace", "fact", "Ace13245 is an enemy", None,
         "Noted player relationship."),
        ("sudo forget my base", "delete", "base", None, None, None,
         "Deleted 'base' memory."),
        ("sudo save this as my mine", "write", "mine", "location", {"x": 30, "y": 11, "z": -80}, None,
         "Saved mining location."),
        ("sudo tp me to my mine", "read", "mine", None, None, "tp {p} 30 11 -80",
         "Read 'mine' from memory. Teleporting to Y=11."),
    ]

    for prompt, action, key, mtype, value, cmd_tmpl, reasoning in MEMORY_OPS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {prompt}")]

        if action == "write":
            msgs.append(tool_call("memory.write", {"player": player, "type": mtype, "key": key, "value": value}))
            msgs.append(tool_result({"ok": True, "key": key}))
            resp = {"risk_level": 4, "commands": [], "reasoning": reasoning}
        elif action == "read":
            args = {"player": player}
            if key:
                args["key"] = key
            msgs.append(tool_call("memory.read", args))
            if key:
                mem_val = value or {"x": random.randint(-500,500), "y": random.randint(60,120), "z": random.randint(-500,500)}
                msgs.append(tool_result({"memories": [{"key": key, "type": mtype or "location", "value": mem_val}]}))
            else:
                msgs.append(tool_result({"memories": [
                    {"key": "home", "type": "location", "value": {"x": 100, "y": 64, "z": -200}},
                    {"key": "favorite_tool", "type": "preference", "value": "diamond pickaxe"},
                ]}))

            if cmd_tmpl:
                cmd = cmd_tmpl.replace("{p}", player)
                result = rcon.command(cmd)
                msgs.append(tool_call("rcon.execute", {"command": cmd}))
                msgs.append(tool_result({"success": not is_error(result), "result": result[:200]}))
                resp = {"risk_level": 3, "commands": [cmd], "reasoning": reasoning}
            else:
                resp = {"risk_level": 4, "commands": [], "reasoning": reasoning}
        elif action == "delete":
            msgs.append(tool_call("memory.read", {"player": player, "key": key}))
            msgs.append(tool_result({"memories": [{"key": key, "type": "location", "value": {"x": -500, "y": 72, "z": 300}}]}))
            resp = {"risk_level": 3, "commands": [], "reasoning": reasoning}

        msgs.append(final_response(resp))
        examples.append({
            "id": f"rcon-memory-{len(examples):03d}",
            "source": "rcon_validated",
            "type": f"memory_{action}",
            "messages": msgs,
        })

    return examples


def gen_wiki_examples(rcon):
    """Wiki lookup → command execution with real RCON results."""
    examples = []

    WIKI = [
        ("sudo best pickaxe enchantments", "pickaxe enchantments 1.21",
         "Pickaxe enchantments:\n- Efficiency V: faster mining\n- Fortune III: more drops from ores\n- Silk Touch: mine blocks directly\n- Unbreaking III: durability\n- Mending: XP repair\nFortune and Silk Touch are mutually exclusive.",
         ["give {p} minecraft:netherite_pickaxe[enchantments={efficiency:5,fortune:3,unbreaking:3,mending:1}] 1"],
         "Max pickaxe with fortune. Fortune and silk touch are exclusive — chose fortune for ore farming."),
        ("sudo what enchants for boots?", "boot enchantments 1.21",
         "Boot enchantments:\n- Protection IV or Blast/Fire/Projectile Protection IV\n- Feather Falling IV: reduces fall damage\n- Depth Strider III or Frost Walker II (exclusive)\n- Soul Speed III: faster on soul sand\n- Swift Sneak III: faster sneaking\n- Unbreaking III, Mending I",
         ["give {p} minecraft:netherite_boots[enchantments={protection:4,feather_falling:4,depth_strider:3,soul_speed:3,unbreaking:3,mending:1}] 1"],
         "All compatible boot enchants. Depth strider over frost walker for general use."),
        ("sudo how do I use the /place command?", "place command 1.21",
         "/place feature <feature> [pos]\n/place jigsaw <pool> <target> <max_depth> [pos]\n/place structure <structure> [pos]\n/place template <template> [pos] [rotation] [mirror]\n\nUseful for placing structures, trees, etc.",
         ["place feature minecraft:oak ~ ~ ~5"],
         "Placed an oak tree feature 5 blocks away using /place."),
        ("sudo what's the command for teams?", "team command minecraft",
         "Team commands:\n/team add <name> [display_name]\n/team join <team> [members]\n/team leave [members]\n/team modify <team> <option> <value>\nOptions: color, friendlyFire, seeFriendlyInvisibles, nametagVisibility, collisionRule, prefix, suffix",
         ["team add red Red Team", "team modify red color red"],
         "Created 'Red Team' with red name color."),
        ("sudo how to make colored text in chat?", "tellraw color formatting minecraft",
         "tellraw @a {\"text\":\"Hello\",\"color\":\"gold\",\"bold\":true}\nColors: black, dark_blue, dark_green, dark_aqua, dark_red, dark_purple, gold, gray, dark_gray, blue, green, aqua, red, light_purple, yellow, white\nFormatting: bold, italic, underlined, strikethrough, obfuscated",
         ['tellraw @a {"text":"Welcome to the server!","color":"gold","bold":true}'],
         "Used tellraw with gold color and bold formatting."),
    ]

    for prompt, query, content, cmds_tmpl, reasoning in WIKI:
        player = random.choice(PLAYERS)
        cmds = [c.replace("{p}", player) for c in cmds_tmpl]
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {prompt}")]

        msgs.append(tool_call("minecraft.wiki_lookup", {"query": query}))
        msgs.append(tool_result({"content": content, "url": f"https://minecraft.wiki/w/{query.split()[0]}"}))

        for cmd in cmds:
            result = rcon.command(cmd)
            msgs.append(tool_call("rcon.execute", {"command": cmd}))
            msgs.append(tool_result({"success": not is_error(result), "result": result[:300]}))

        resp = {"risk_level": 3, "commands": cmds, "reasoning": reasoning}
        msgs.append(final_response(resp))

        examples.append({
            "id": f"rcon-wiki-{len(examples):03d}",
            "source": "rcon_validated",
            "type": "wiki_lookup",
            "messages": msgs,
        })

    return examples


def gen_error_correction_examples(rcon):
    """Commands that fail, get corrected, and succeed — with real RCON errors."""
    examples = []

    ERRORS = [
        ("sudo give me a bed", "give {p} minecraft:bed 1", "give {p} minecraft:white_bed 1",
         "No item 'bed' — correct ID is 'white_bed' in 1.21."),
        ("sudo give me cooked beef", "give {p} minecraft:steak 1", "give {p} minecraft:cooked_beef 1",
         "Item is 'cooked_beef', not 'steak'."),
        ("sudo give me speed", "effect give {p} minecraft:speed", "effect give {p} minecraft:speed 300 1",
         "Effect needs duration and amplifier arguments."),
        ("sudo give me a log", "give {p} minecraft:log 64", "give {p} minecraft:oak_log 64",
         "No item 'log' — must specify wood type: oak_log, spruce_log, etc."),
        ("sudo fill with stone 10", "fill ~ ~ ~ ~10 ~10 ~10 minecraft:stone 10", "fill ~ ~ ~ ~10 ~10 ~10 minecraft:stone",
         "Fill doesn't take a trailing count number."),
    ]

    for prompt, wrong_cmd_tmpl, right_cmd_tmpl, reasoning in ERRORS:
        player = random.choice(PLAYERS)
        wrong_cmd = wrong_cmd_tmpl.replace("{p}", player)
        right_cmd = right_cmd_tmpl.replace("{p}", player)
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {prompt}")]

        # Wrong attempt
        wrong_result = rcon.command(wrong_cmd)
        msgs.append(tool_call("rcon.execute", {"command": wrong_cmd}))
        msgs.append(tool_result({"success": False, "result": wrong_result[:300]}))

        # Corrected attempt
        right_result = rcon.command(right_cmd)
        msgs.append(tool_call("rcon.execute", {"command": right_cmd}))
        msgs.append(tool_result({"success": not is_error(right_result), "result": right_result[:300]}))

        resp = {"risk_level": 3, "commands": [right_cmd], "reasoning": f"First attempt failed: {wrong_result[:100]}. {reasoning}"}
        msgs.append(final_response(resp))

        examples.append({
            "id": f"rcon-errfix-{len(examples):03d}",
            "source": "rcon_validated",
            "type": "error_correction",
            "messages": msgs,
        })

    return examples


def gen_chained_examples(rcon):
    """Multi-tool chained examples with real RCON."""
    examples = []

    CHAINS = [
        ("sudo gear me up for the nether", [
            ("rcon.execute", {"command": "give {p} minecraft:netherite_sword[enchantments={sharpness:5,fire_aspect:2,unbreaking:3,mending:1}] 1"}),
            ("rcon.execute", {"command": "give {p} minecraft:netherite_pickaxe[enchantments={efficiency:5,fortune:3,unbreaking:3,mending:1}] 1"}),
            ("rcon.execute", {"command": "give {p} minecraft:netherite_chestplate[enchantments={protection:4,unbreaking:3,mending:1}] 1"}),
            ("rcon.execute", {"command": "effect give {p} minecraft:fire_resistance 6000 0"}),
            ("rcon.execute", {"command": "give {p} minecraft:golden_apple 16"}),
            ("rcon.execute", {"command": "give {p} minecraft:ender_pearl 16"}),
        ], "Full nether loadout: netherite sword+pick, protection armor, fire resistance, golden apples, pearls."),
        ("sudo prepare me for an end fight", [
            ("rcon.execute", {"command": "give {p} minecraft:netherite_sword[enchantments={sharpness:5,unbreaking:3,mending:1}] 1"}),
            ("rcon.execute", {"command": "give {p} minecraft:bow[enchantments={power:5,infinity:1,unbreaking:3}] 1"}),
            ("rcon.execute", {"command": "give {p} minecraft:arrow 64"}),
            ("rcon.execute", {"command": "give {p} minecraft:golden_apple 32"}),
            ("rcon.execute", {"command": "give {p} minecraft:ender_pearl 16"}),
            ("rcon.execute", {"command": "effect give {p} minecraft:slow_falling 600 0"}),
            ("rcon.execute", {"command": "give {p} minecraft:cobblestone 128"}),
        ], "End fight kit: weapons, slow falling for void safety, blocks for pillaring, pearls for dodging."),
    ]

    for prompt, tool_calls, reasoning in CHAINS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {prompt}")]

        all_cmds = []
        for tool_name, args_tmpl in tool_calls:
            args = {}
            for k, v in args_tmpl.items():
                args[k] = v.replace("{p}", player) if isinstance(v, str) else v

            if tool_name == "rcon.execute":
                cmd = args["command"]
                result = rcon.command(cmd)
                msgs.append(tool_call(tool_name, args))
                msgs.append(tool_result({"success": not is_error(result), "result": result[:300]}))
                all_cmds.append(cmd)
            else:
                msgs.append(tool_call(tool_name, args))
                msgs.append(tool_result({"ok": True}))

        resp = {"risk_level": 3, "commands": all_cmds, "reasoning": reasoning}
        msgs.append(final_response(resp))

        examples.append({
            "id": f"rcon-chain-{len(examples):03d}",
            "source": "rcon_validated",
            "type": "chained",
            "messages": msgs,
        })

    return examples


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--rcon-host", default="192.168.0.244")
    parser.add_argument("--rcon-port", type=int, default=25577)
    parser.add_argument("--rcon-pass", default="REDACTED_RCON")
    args = parser.parse_args()

    print(f"Connecting to RCON at {args.rcon_host}:{args.rcon_port}...")
    rcon = SimpleRCON(args.rcon_host, args.rcon_port, args.rcon_pass)
    print("Connected.\n")

    all_examples = []

    print("Generating script examples...")
    scripts = gen_script_examples(rcon)
    all_examples.extend(scripts)
    print(f"  {len(scripts)} script examples")

    print("Generating direct command examples...")
    direct = gen_direct_command_examples(rcon)
    all_examples.extend(direct)
    print(f"  {len(direct)} direct command examples")

    print("Generating memory examples...")
    memory = gen_memory_examples(rcon)
    all_examples.extend(memory)
    print(f"  {len(memory)} memory examples")

    print("Generating wiki examples...")
    wiki = gen_wiki_examples(rcon)
    all_examples.extend(wiki)
    print(f"  {len(wiki)} wiki examples")

    print("Generating error correction examples...")
    errors = gen_error_correction_examples(rcon)
    all_examples.extend(errors)
    print(f"  {len(errors)} error correction examples")

    print("Generating chained multi-tool examples...")
    chained = gen_chained_examples(rcon)
    all_examples.extend(chained)
    print(f"  {len(chained)} chained examples")

    # Stats
    success_count = sum(1 for e in all_examples if e.get("all_success", True))
    print(f"\nTotal: {len(all_examples)} examples")
    print(f"  RCON validated: {sum(1 for e in all_examples if e.get('rcon_validated', False))}")

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

    print(f"Written to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()