#!/usr/bin/env python3
"""
generate_new_tool_examples.py — Generate training examples for under-represented tools.

Covers: world.render, server.config, training.save, user.ask, and additional
world.scan_area + world.redstone_trace patterns.

Output: data/raw/new_tool_training.jsonl
"""

import json
import random
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent.parent
OUTPUT = ROOT / "data" / "raw" / "new_tool_training.jsonl"

SYSTEM_PROMPT = """/no_think
You are Mortdecai, a Minecraft 1.21 AI for a Paper server with FAWE, WorldGuard, CoreProtect, EssentialsX, Vault, LuckPerms.

TOOLS: rcon.execute, minecraft.lookup, plugin.docs_lookup, world.player_info, world.server_state, world.nearby_entities, world.scan_area, world.redstone_trace, world.render, server.config, memory.read, memory.write, journal.read, journal.write, log.query, user.ask, script.write, script.validate, script.execute, script.read, script.list, script.delete, script.schedule, training.save

To call a tool, respond with:
<tool_call>
{"name": "tool_name", "arguments": {...}}
</tool_call>

After all tool calls, respond with final JSON: {"commands": [...], "message": "...", "reasoning": "..."}"""

PLAYERS = ["slingshooter08", "SwiftWolf", "DarkWolf", "BraveWolf", "CraftMaster99", "RedstoneWiz"]


def tc(name, args):
    """Format a tool_call block."""
    return f'<tool_call>\n{json.dumps({"name": name, "arguments": args})}\n</tool_call>'


def tr(result):
    """Format a tool result."""
    return json.dumps(result, ensure_ascii=True)


def final(commands=None, message=None, reasoning=""):
    return json.dumps({"commands": commands or [], "message": message, "reasoning": reasoning})


# ── world.render examples ──────────────────────────────────────────

def render_build_inspection():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo show me what my house looks like from above"},
        {"role": "assistant", "content": tc("world.player_info", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "player": p, "position": {"x": 50, "y": 68, "z": -30}, "gamemode": "survival"})},
        {"role": "assistant", "content": tc("world.render", {"x1": 40, "y1": 60, "z1": -40, "x2": 60, "y2": 80, "z2": -20, "mode": "top_down", "y_level": 68})},
        {"role": "tool", "content": tr({"ok": True, "render": "....WWWWWW....\n...W......W...\n..W..DD..W...\n..W......W...\n..W......W...\n...WWWWWW....\n\nLegend: W=oak_planks D=oak_door .=air", "mode": "top_down", "blocks_rendered": 24})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f'tellraw {p} {{"text":"Here is a top-down view of your house at Y=68. It is a rectangular oak plank structure with a door on the south side.","color":"green"}}'})},
        {"role": "tool", "content": tr({"success": True, "result": ""})},
        {"role": "assistant", "content": final([], None, "Rendered player's house from top-down view at their Y level. Communicated the layout via tellraw.")},
    ]}


def render_cross_section():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo can you see how deep my mine goes?"},
        {"role": "assistant", "content": tc("world.player_info", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "player": p, "position": {"x": -100, "y": 20, "z": 200}, "gamemode": "survival"})},
        {"role": "assistant", "content": tc("world.render", {"x1": -100, "y1": -60, "z1": 195, "x2": -100, "y2": 68, "z2": 205, "mode": "cross_z"})},
        {"role": "tool", "content": tr({"ok": True, "render": "SSSSSSSSS\nSSSS.SSSS\nSSS...SSS\nSS.....SS\nS.......S\n.........\n.T....T..\n.........\nSSSSSSSSS\n\nLegend: S=stone .=air T=torch", "mode": "cross_z", "blocks_rendered": 45})},
        {"role": "assistant", "content": final([], None, "Rendered cross-section of player's mine. It extends from Y=68 surface down approximately 40 blocks with torches placed along the shaft.")},
    ]}


def render_redstone_layout():
    p = "RedstoneWiz"
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo render my redstone circuit please"},
        {"role": "assistant", "content": tc("world.player_info", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "player": p, "position": {"x": 10, "y": 64, "z": 10}, "gamemode": "creative"})},
        {"role": "assistant", "content": tc("world.scan_area", {"center": {"x": 10, "y": 64, "z": 10}, "radius": 15, "filter": "redstone"})},
        {"role": "tool", "content": tr({"ok": True, "blocks": [{"pos": [5, 64, 10], "type": "lever"}, {"pos": [6, 64, 10], "type": "redstone_wire"}, {"pos": [7, 64, 10], "type": "redstone_wire"}, {"pos": [8, 64, 10], "type": "repeater"}, {"pos": [9, 64, 10], "type": "redstone_wire"}, {"pos": [10, 64, 10], "type": "redstone_lamp"}], "count": 6})},
        {"role": "assistant", "content": tc("world.render", {"x1": 3, "y1": 63, "z1": 8, "x2": 12, "y2": 65, "z2": 12, "mode": "top_down", "y_level": 64})},
        {"role": "tool", "content": tr({"ok": True, "render": "...........\n..L-RR>R-O.\n...........\n\nLegend: L=lever -=redstone_wire >=repeater O=lamp .=air", "mode": "top_down", "blocks_rendered": 6})},
        {"role": "assistant", "content": final([], None, "Scanned and rendered redstone circuit: lever → 2 wire → repeater → 1 wire → lamp. Linear circuit running east at Y=64.")},
    ]}


# ── server.config examples ──────────────────────────────────────────

def config_read_gamerules():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo what are the current gamerules?"},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "keepInventory"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "keepInventory", "value": "false"})},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "mobGriefing"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "mobGriefing", "value": "true"})},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "doDaylightCycle"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "doDaylightCycle", "value": "true"})},
        {"role": "assistant", "content": final([], None, "Checked key gamerules: keepInventory=false (items drop on death), mobGriefing=true (mobs can damage terrain), doDaylightCycle=true (time passes normally).")},
    ]}


def config_set_keepinventory():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo turn on keep inventory please"},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "keepInventory"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "keepInventory", "value": "false"})},
        {"role": "assistant", "content": tc("server.config", {"action": "set", "setting": "keepInventory", "value": "true"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "keepInventory", "value": "true", "previous": "false"})},
        {"role": "assistant", "content": final(["gamerule keepInventory true"], None, "Checked keepInventory was false, then enabled it. Players will now keep items on death.")},
    ]}


def config_difficulty():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo make the game harder"},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "difficulty"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "difficulty", "value": "normal"})},
        {"role": "assistant", "content": tc("server.config", {"action": "set", "setting": "difficulty", "value": "hard"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "difficulty", "value": "hard", "previous": "normal"})},
        {"role": "assistant", "content": final(["difficulty hard"], None, "Changed difficulty from normal to hard. Zombies can break doors, hunger can kill, and mobs deal more damage.")},
    ]}


def config_worldborder():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo set the world border to 5000 blocks"},
        {"role": "assistant", "content": tc("server.config", {"action": "read", "setting": "worldborder"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "worldborder", "value": "60000000"})},
        {"role": "assistant", "content": tc("server.config", {"action": "set", "setting": "worldborder", "value": "5000"})},
        {"role": "tool", "content": tr({"ok": True, "setting": "worldborder", "value": "5000", "previous": "60000000"})},
        {"role": "assistant", "content": final(["worldborder set 5000"], None, "Set world border from default (60M) to 5000 blocks diameter. Players cannot travel beyond 2500 blocks from center.")},
    ]}


# ── user.ask examples ──────────────────────────────────────────────

def user_ask_ambiguous_kill():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo kill that guy"},
        {"role": "assistant", "content": tc("world.nearby_entities", {"player": p, "radius": 30})},
        {"role": "tool", "content": tr({"ok": True, "entities": [{"type": "player", "name": "SwiftWolf", "distance": 8}, {"type": "player", "name": "DarkWolf", "distance": 15}, {"type": "zombie", "count": 3, "nearest_distance": 12}]})},
        {"role": "assistant", "content": tc("user.ask", {"player": p, "question": "There are 2 players nearby: SwiftWolf (8 blocks) and DarkWolf (15 blocks), plus 3 zombies. Who did you mean by 'that guy'?"})},
        {"role": "tool", "content": tr({"ok": True, "response": "the zombies not the players lol"})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": "kill @e[type=minecraft:zombie,distance=..30]"})},
        {"role": "tool", "content": tr({"success": True, "result": "Killed 3 entities"})},
        {"role": "assistant", "content": final(["kill @e[type=minecraft:zombie,distance=..30]"], None, "Ambiguous 'kill that guy' — checked nearby entities, found both players and zombies. Asked for clarification. Player meant the zombies. Killed all zombies within 30 blocks.")},
    ]}


def user_ask_destructive_fill():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo fill this whole area with lava"},
        {"role": "assistant", "content": tc("world.player_info", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "player": p, "position": {"x": 30, "y": 65, "z": -20}, "gamemode": "survival"})},
        {"role": "assistant", "content": tc("world.nearby_entities", {"player": p, "radius": 50})},
        {"role": "tool", "content": tr({"ok": True, "entities": [{"type": "player", "name": "BraveWolf", "distance": 20}, {"type": "cow", "count": 5, "nearest_distance": 10}]})},
        {"role": "assistant", "content": tc("user.ask", {"player": p, "question": "Filling with lava is destructive and BraveWolf is 20 blocks away. How large an area? And should I warn BraveWolf first?"})},
        {"role": "tool", "content": tr({"ok": True, "response": "just a 5x5 pit right here, and yeah warn them"})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f'tellraw BraveWolf {{"text":"Warning: {p} is filling a small area with lava near you. Stay clear!","color":"red"}}'})},
        {"role": "tool", "content": tr({"success": True, "result": ""})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": "fill 28 60 -22 32 64 -18 minecraft:lava"})},
        {"role": "tool", "content": tr({"success": True, "result": "Successfully filled 125 blocks"})},
        {"role": "assistant", "content": final(["fill 28 60 -22 32 64 -18 minecraft:lava"], None, "Destructive lava fill — asked for clarification on scope. Player wanted 5x5 pit. Warned nearby player BraveWolf first, then filled.")},
    ]}


def user_ask_ban_request():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo ban SwiftWolf he was griefing"},
        {"role": "assistant", "content": tc("log.query", {"type": "actions", "player": "SwiftWolf", "limit": 10})},
        {"role": "tool", "content": tr({"ok": True, "results": [{"type": "action", "player": "SwiftWolf", "detail": "placed TNT at (100, 65, 200)", "age": "3m ago"}, {"type": "action", "player": "SwiftWolf", "detail": "broke oak_planks at (102, 66, 201)", "age": "4m ago"}], "count": 2})},
        {"role": "assistant", "content": tc("user.ask", {"player": p, "question": "I see SwiftWolf placed TNT and broke some oak planks recently. Are you sure you want a permanent ban? Other options: tempban (hours), kick, or jail."})},
        {"role": "tool", "content": tr({"ok": True, "response": "just kick him for now and warn him"})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": "kick SwiftWolf Griefing warning: Do not place TNT near player builds."})},
        {"role": "tool", "content": tr({"success": True, "result": "Kicked SwiftWolf"})},
        {"role": "assistant", "content": final(["kick SwiftWolf Griefing warning: Do not place TNT near player builds."], None, "Investigated griefing claim — logs confirmed TNT placement and block breaking. Asked about severity preference. Player chose kick with warning instead of ban.")},
    ]}


# ── training.save examples ──────────────────────────────────────────

def training_save_syntax_discovery():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo give me a colored leather helmet"},
        {"role": "assistant", "content": tc("minecraft.lookup", {"query": "leather armor color dye command 1.21"})},
        {"role": "tool", "content": tr({"ok": True, "results": [{"title": "Dyed Armor", "text": "In 1.21, use: give @p leather_helmet[dyed_color=16711680] for red. Color is decimal RGB."}]})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f"give {p} minecraft:leather_helmet[dyed_color=16711680] 1"})},
        {"role": "tool", "content": tr({"success": False, "result": "Invalid or unknown argument"})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f"give {p} minecraft:leather_helmet 1"})},
        {"role": "tool", "content": tr({"success": True, "result": "Gave 1 [Leather Cap]"})},
        {"role": "assistant", "content": tc("training.save", {"lesson": "Paper RCON rejects dyed_color component syntax in give command. Must give plain leather armor — dyeing requires a crafting table with dyes.", "category": "syntax_discovery", "input_summary": "give colored leather helmet", "correct_output": "give @p minecraft:leather_helmet 1", "wrong_output": "give @p minecraft:leather_helmet[dyed_color=16711680] 1", "confidence": 0.9})},
        {"role": "tool", "content": tr({"ok": True, "saved": True, "id": "lesson-042"})},
        {"role": "assistant", "content": final([f"give {p} minecraft:leather_helmet 1"], None, "Attempted colored leather helmet via component syntax — Paper RCON rejected it. Gave plain helmet instead. Saved lesson about RCON component limitations.")},
    ]}


def training_save_error_correction():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo set up a repeating command block"},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f"give {p} minecraft:repeating_command_block 1"})},
        {"role": "tool", "content": tr({"success": False, "result": "Unknown item: minecraft:repeating_command_block"})},
        {"role": "assistant", "content": tc("minecraft.lookup", {"query": "repeating command block item id 1.21"})},
        {"role": "tool", "content": tr({"ok": True, "results": [{"title": "Command Block", "text": "In 1.21: minecraft:command_block (impulse), minecraft:chain_command_block, minecraft:repeating_command_block. Note: give may not work for all variants on Paper."}]})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f"give {p} minecraft:command_block 1"})},
        {"role": "tool", "content": tr({"success": True, "result": "Gave 1 [Command Block]"})},
        {"role": "assistant", "content": tc("training.save", {"lesson": "Paper RCON may reject 'give minecraft:repeating_command_block' — use 'give minecraft:command_block' and let the player change the type in-game by right-clicking.", "category": "error_correction", "input_summary": "give repeating command block", "correct_output": "give @p minecraft:command_block 1", "wrong_output": "give @p minecraft:repeating_command_block 1", "confidence": 0.8})},
        {"role": "tool", "content": tr({"ok": True, "saved": True, "id": "lesson-043"})},
        {"role": "assistant", "content": final([f"give {p} minecraft:command_block 1"], None, "Repeating command block give failed. Fell back to regular command block. Player can change type in-game. Saved error correction lesson.")},
    ]}


def training_save_redstone_rule():
    p = "RedstoneWiz"
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo my repeater isn't passing signal, what's wrong?"},
        {"role": "assistant", "content": tc("world.player_info", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "player": p, "position": {"x": 15, "y": 64, "z": 25}, "gamemode": "creative"})},
        {"role": "assistant", "content": tc("world.redstone_trace", {"start": {"x": 15, "y": 64, "z": 25}})},
        {"role": "tool", "content": tr({"ok": True, "trace": [{"pos": [13, 64, 25], "type": "lever", "powered": True}, {"pos": [14, 64, 25], "type": "wire", "power": 14}, {"pos": [15, 64, 25], "type": "repeater", "facing": "west", "powered": False}], "break_point": {"pos": [15, 64, 25], "reason": "repeater facing wrong direction — input is from east but repeater faces west (output goes west, expects input from east side but signal arrives on wrong side)"}})},
        {"role": "assistant", "content": tc("rcon.execute", {"command": f'tellraw {p} {{"text":"Found it! Your repeater at (15,64,25) faces west but signal comes from the east side. The repeater input/output are swapped. Break and re-place it facing east.","color":"yellow"}}'})},
        {"role": "tool", "content": tr({"success": True, "result": ""})},
        {"role": "assistant", "content": tc("training.save", {"lesson": "Repeater facing=west means it outputs west and expects input from east. If signal comes from west side, the repeater won't fire — it needs to be rotated.", "category": "redstone_rule", "input_summary": "repeater not passing signal", "correct_output": "Repeater must face the direction signal should travel (output direction)", "confidence": 0.95})},
        {"role": "tool", "content": tr({"ok": True, "saved": True, "id": "lesson-044"})},
        {"role": "assistant", "content": final([], None, "Traced redstone circuit and found repeater facing wrong direction. Explained fix via tellraw. Saved redstone rule about repeater facing direction.")},
    ]}


# ── Additional log.query + user.ask combos ──────────────────────────

def log_query_death_investigation():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: pray why do I keep dying?"},
        {"role": "assistant", "content": tc("log.query", {"type": "deaths", "player": p, "limit": 5})},
        {"role": "tool", "content": tr({"ok": True, "results": [{"type": "death", "player": p, "detail": "was slain by Zombie", "age": "2m ago"}, {"type": "death", "player": p, "detail": "was slain by Zombie", "age": "8m ago"}, {"type": "death", "player": p, "detail": "fell from a high place", "age": "15m ago"}, {"type": "death", "player": p, "detail": "was slain by Skeleton", "age": "20m ago"}], "count": 4})},
        {"role": "assistant", "content": tc("journal.read", {"player": p})},
        {"role": "tool", "content": tr({"ok": True, "journal": "Newer player, tends to explore caves without armor. Dies frequently to mobs."})},
        {"role": "assistant", "content": final(
            [f"give {p} minecraft:iron_chestplate 1", f"give {p} minecraft:iron_sword 1", f"effect give {p} minecraft:regeneration 120 1"],
            "I have peered into the records of your mortal misfortunes. FOUR deaths in twenty minutes! Twice by zombie, once by skeleton, and a tragic fall. The pattern is clear: you venture forth without proper protection. Here — armor and a blade. Stop dying.",
            "Queried death logs: 4 recent deaths, mostly mob-related. Journal confirms player is under-equipped. Providing armor, weapon, and regen.")},
    ]}


def log_query_who_was_here():
    p = random.choice(PLAYERS)
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Player {p}: sudo who was online earlier today?"},
        {"role": "assistant", "content": tc("log.query", {"type": "joins", "limit": 10})},
        {"role": "tool", "content": tr({"ok": True, "results": [
            {"type": "join", "player": "SwiftWolf", "detail": "joined the game", "age": "15m ago"},
            {"type": "join", "player": "DarkWolf", "detail": "left the game", "age": "30m ago"},
            {"type": "join", "player": "DarkWolf", "detail": "joined the game", "age": "2h ago"},
            {"type": "join", "player": "BraveWolf", "detail": "joined the game", "age": "3h ago"},
            {"type": "join", "player": "BraveWolf", "detail": "left the game", "age": "1h ago"},
        ], "count": 5})},
        {"role": "assistant", "content": final([], None, "Join/leave log shows: SwiftWolf joined 15m ago (still on), DarkWolf was on for ~1.5h (left 30m ago), BraveWolf played for 2h (left 1h ago).")},
    ]}


def main():
    random.seed(42)

    generators = [
        # world.render (3 scenarios × 2 variations)
        ("world.render: build inspection", render_build_inspection, 2),
        ("world.render: cross section", render_cross_section, 2),
        ("world.render: redstone layout", render_redstone_layout, 2),
        # server.config (4 scenarios × 2 variations)
        ("server.config: read gamerules", config_read_gamerules, 2),
        ("server.config: set keepInventory", config_set_keepinventory, 2),
        ("server.config: difficulty change", config_difficulty, 2),
        ("server.config: worldborder", config_worldborder, 2),
        # user.ask (3 scenarios × 2 variations)
        ("user.ask: ambiguous kill", user_ask_ambiguous_kill, 2),
        ("user.ask: destructive fill", user_ask_destructive_fill, 2),
        ("user.ask: ban request", user_ask_ban_request, 2),
        # training.save (3 scenarios × 2 variations)
        ("training.save: syntax discovery", training_save_syntax_discovery, 2),
        ("training.save: error correction", training_save_error_correction, 2),
        ("training.save: redstone rule", training_save_redstone_rule, 2),
        # log.query combos (2 scenarios × 2 variations)
        ("log.query: death investigation", log_query_death_investigation, 2),
        ("log.query: who was here", log_query_who_was_here, 2),
    ]

    all_examples = []
    for name, gen_fn, count in generators:
        for _ in range(count):
            all_examples.append(gen_fn())
        print(f"  {name}: {count}")

    random.shuffle(all_examples)

    with open(OUTPUT, "w") as f:
        for ex in all_examples:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    print(f"\nTotal: {len(all_examples)} examples written to {OUTPUT}")


if __name__ == "__main__":
    main()