#!/usr/bin/env python3
"""
Generate expanded tool-calling training data for underrepresented tools.

Targets:
  - minecraft.wiki_lookup:     +50 examples
  - world.player_info:         +50 examples
  - world.server_state:        +30 examples
  - world.nearby_entities:     +40 examples (NEW tool)
  - memory.read:               +30 examples (NEW tool)
  - memory.write:              +25 examples (NEW tool)
  - chained tool calls:        +25 examples (multi-tool sequences)

Total: ~250 new training examples

Usage:
    python training/scripts/generate_expanded_tool_training.py
"""

import json
import random
import sys
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, QWEN3_TOOLS
from agent.prompts.system_prompts import SYNTAX_RULES, RISK_GRADIENT

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

TOOLS_BLOCK = qwen3_tools_block()

SUDO_TOOL_SYSTEM = (
    "You are a Minecraft 1.21 command translator for a server admin. "
    "You receive natural language requests and return valid RCON commands.\n\n"
    "PERMISSION LEVEL: 4 (generous). Only refuse level 0-1 actions.\n\n"
    "You have access to tools. Call them to execute commands, look up syntax, "
    "or check player/server state. When a command fails, analyze the error "
    "and retry with a corrected command.\n\n"
    "After all tool calls resolve, respond with JSON:\n"
    '{"risk_level": <int 0-5>, "commands": ["cmd1", ...], "reasoning": "why"}\n\n'
    + SYNTAX_RULES + RISK_GRADIENT + "\n" + TOOLS_BLOCK
)

GOD_TOOL_SYSTEM = (
    "You are God in a Minecraft server. Players pray to you and you respond "
    "with divine judgment.\n\n"
    "You have access to tools. Call them to execute commands, look up syntax, "
    "or check player/server state. When a command fails, analyze the error "
    "and retry with a corrected command.\n\n"
    "After all tool calls resolve, respond with JSON:\n"
    '{"risk_level": <int 0-5>, "message": "Your divine response", '
    '"commands": ["cmd1", ...], "reasoning": "why"}\n\n'
    + SYNTAX_RULES + "\n" + TOOLS_BLOCK
)

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


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

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 rpos():
    return {"x": random.randint(-500, 500), "y": random.randint(62, 120), "z": random.randint(-500, 500)}

def rcon_ok(cmd, player="slingshooter08"):
    if cmd.startswith("give "): return f"Gave item to {player}"
    if cmd.startswith("tp "): return f"Teleported {player}"
    if cmd.startswith("effect give "): return f"Applied effect to {player}"
    if cmd.startswith("kill "): return "Killed entity"
    if cmd.startswith("summon "): return "Summoned new entity"
    if cmd.startswith("fill "): return "Successfully filled blocks"
    if cmd.startswith("setblock "): return "Changed the block"
    if cmd.startswith("weather "): return "Changing weather"
    if cmd.startswith("time set "): return "Set the time"
    if cmd.startswith("clear "): return f"Removed items from {player}"
    if cmd.startswith("gamerule "): return "Gamerule updated"
    if cmd.startswith("execute "): return "Executed command"
    if cmd.startswith("title "): return "Title displayed"
    if cmd.startswith("playsound "): return f"Played sound to {player}"
    if cmd.startswith("scoreboard "): return "Scoreboard updated"
    return "Command executed successfully"

def make_example(eid, etype, messages):
    return {"id": eid, "source": "expanded_tool_training", "type": etype, "messages": messages}


# ── Wiki Lookup Examples ──────────────────────────────────────────────────

WIKI_EXAMPLES = [
    ("sudo give me a mace with the best enchantments", "mace enchantments 1.21",
     "Mace enchantments in 1.21:\n- Density (1-5): Increases damage based on fall distance\n- Breach (1-4): Reduces armor effectiveness\n- Wind Burst (1-3): Launches attacker upward on hit\n- Smite (1-5): Extra damage to undead\n- Bane of Arthropods (1-5): Extra damage to arthropods\n- Fire Aspect (1-2): Sets target on fire\n- Unbreaking (1-3): Durability\n- Mending (1): XP repair\nNote: Density, Breach, Smite, and Bane of Arthropods are mutually exclusive.",
     "https://minecraft.wiki/w/Mace", ["give {p} minecraft:mace[enchantments={{density:5,wind_burst:3,fire_aspect:2,unbreaking:3,mending:1}}] 1"],
     "Looked up mace enchantments. Density and breach are exclusive — chose density for max fall damage. Added wind_burst for combos."),

    ("sudo what are all the copper block variants?", "copper block variants oxidation 1.21",
     "Copper blocks in 1.21 have 4 oxidation states and 2 waxing states:\n- copper_block, exposed_copper, weathered_copper, oxidized_copper\n- Cut variants: cut_copper, exposed_cut_copper, weathered_cut_copper, oxidized_cut_copper\n- Slab/stair variants for each\n- Waxed variants of all (prefix waxed_)\n- Copper bulb, copper door, copper trapdoor, copper grate (also with oxidation states)\nUse honeycomb to wax. Axe to scrape oxidation.",
     "https://minecraft.wiki/w/Copper_Block", ["give {p} minecraft:copper_block 64", "give {p} minecraft:exposed_copper 64", "give {p} minecraft:weathered_copper 64", "give {p} minecraft:oxidized_copper 64"],
     "Looked up copper variants. Gave one stack of each oxidation state."),

    ("sudo how do I make a command block chain?", "command block chain repeating impulse 1.21",
     "Command block types:\n- Impulse (orange): Executes once when triggered\n- Chain (green): Executes when the block behind it activates\n- Repeating (purple): Executes every tick while powered\n\nPlacement: command_block (impulse), chain_command_block, repeating_command_block\nModes: Always Active or Needs Redstone\nConditional: Only runs if previous block in chain succeeded",
     "https://minecraft.wiki/w/Command_Block", ["give {p} minecraft:command_block 1", "give {p} minecraft:chain_command_block 4", "give {p} minecraft:repeating_command_block 1"],
     "Looked up command block types. Gave impulse, chain, and repeating blocks for a chain setup."),

    ("sudo what's the syntax for the attribute command?", "attribute command syntax 1.21",
     "attribute <target> <attribute> base get|set <value>\nattribute <target> <attribute> modifier add <uuid> <name> <value> add|multiply|multiply_base\nattribute <target> <attribute> modifier remove <uuid>\n\nCommon attributes:\n- minecraft:generic.max_health (default 20)\n- minecraft:generic.movement_speed (default 0.1)\n- minecraft:generic.attack_damage (default 1)\n- minecraft:generic.armor (default 0)\n- minecraft:generic.knockback_resistance (0-1)",
     "https://minecraft.wiki/w/Attribute", ["attribute {p} minecraft:generic.max_health base set 40"],
     "Looked up attribute syntax. Set max health to 40 (double normal)."),

    ("sudo what does the wind charge do?", "wind charge item minecraft 1.21",
     "Wind Charge: Projectile item added in 1.21\n- Thrown like a snowball\n- On impact: creates wind burst that knocks back entities and activates redstone\n- Players can use it for movement: throw at feet for a boost\n- Breeze mob shoots these naturally\n- Item: minecraft:wind_charge",
     "https://minecraft.wiki/w/Wind_Charge", ["give {p} minecraft:wind_charge 16"],
     "Looked up wind charge mechanics. Gave 16 wind charges."),

    ("sudo how does the trial spawner work?", "trial spawner mechanics 1.21",
     "Trial Spawner: Found in Trial Chambers\n- Spawns mobs based on nearby players\n- Becomes 'ominous' if player has Bad Omen\n- Cannot be obtained with Silk Touch\n- Ejects rewards when all mobs defeated\n- Cooldown of 30 minutes per player\n- Item: minecraft:trial_spawner (creative only)",
     "https://minecraft.wiki/w/Trial_Spawner", ["give {p} minecraft:trial_spawner 1"],
     "Looked up trial spawner. It's a creative-only item. Giving one."),

    ("sudo what are the vault block mechanics?", "vault block trial chambers 1.21",
     "Vault: Loot container in Trial Chambers\n- Each player can open it once with a trial key\n- Drops random loot (armor trims, enchanted books, etc.)\n- Ominous vault: requires ominous trial key, better loot\n- Items: minecraft:vault, minecraft:trial_key, minecraft:ominous_trial_key",
     "https://minecraft.wiki/w/Vault", ["give {p} minecraft:vault 1", "give {p} minecraft:trial_key 3"],
     "Looked up vault mechanics. Gave a vault and 3 trial keys."),

    ("sudo what food gives the most saturation?", "food saturation values minecraft",
     "Highest saturation foods:\n1. Golden carrot: 14.4 saturation, 6 hunger\n2. Cooked porkchop: 12.8 saturation, 8 hunger\n3. Steak (cooked beef): 12.8 saturation, 8 hunger\n4. Cooked mutton: 9.6 saturation, 6 hunger\n5. Cooked salmon: 9.6 saturation, 6 hunger\n6. Suspicious stew varies by flower\nGolden apple: 9.6 sat, enchanted golden apple: 9.6 sat",
     "https://minecraft.wiki/w/Food", ["give {p} minecraft:golden_carrot 64", "give {p} minecraft:cooked_beef 64"],
     "Looked up food saturation. Golden carrots have the best saturation. Gave those plus steak."),

    ("sudo what's the difference between smite and sharpness?", "smite vs sharpness enchantment damage",
     "Sharpness: +0.5 * level + 0.5 damage to all mobs (max V = 3 extra)\nSmite: +2.5 * level damage to undead only (max V = 12.5 extra)\nBane of Arthropods: +2.5 * level to arthropods only\n\nThese three are mutually exclusive.\nSmite is better vs undead (zombies, skeletons, wither, phantoms).\nSharpness is better for general use.",
     "https://minecraft.wiki/w/Sharpness", ["give {p} minecraft:diamond_sword[enchantments={{sharpness:5,unbreaking:3,mending:1,looting:3,sweeping_edge:3}}] 1"],
     "Looked up sharpness vs smite. Sharpness is better for general use. Gave max enchanted sword."),

    ("sudo what's the recipe for a lodestone?", "lodestone recipe crafting minecraft",
     "Lodestone recipe:\n- 8 chiseled stone bricks (surrounding)\n- 1 netherite ingot (center)\nCrafting: 3x3 grid with chiseled stone bricks around netherite ingot\nUsage: Right-click with compass to bind it. Works across dimensions.",
     "https://minecraft.wiki/w/Lodestone", ["give {p} minecraft:lodestone 1"],
     "Looked up lodestone recipe. It requires a netherite ingot. Giving one directly."),

    ("sudo what paintings are available in 1.21?", "paintings list new 1.21",
     "1.21 added 20 new paintings by artist Sarah Boeving:\n- backyard, bouquet, cavebird, changing, cotan, endboss, fern, finding, lowmist, orb, owlemons, passage, pond, sunflowers, tides, unpacked\nPaintings are placed with: give <player> minecraft:painting 1\nThe painting chosen depends on wall size.",
     "https://minecraft.wiki/w/Painting", ["give {p} minecraft:painting 5"],
     "Looked up 1.21 paintings. 20 new ones added. Gave 5 paintings to place."),

    ("sudo how do I use the /damage command?", "damage command syntax 1.21",
     "damage <target> <amount> [<damageType>] [at <location>] [by <entity>] [from <cause>]\n\nExamples:\ndamage @p 5 minecraft:fall\ndamage @e[type=zombie,limit=1] 10 minecraft:player_attack by @p\n\nDamage types: fall, fire, drown, starve, freeze, player_attack, mob_attack, explosion, etc.",
     "https://minecraft.wiki/w/Commands/damage", ["damage @e[type=zombie,distance=..10] 20 minecraft:player_attack"],
     "Looked up /damage syntax. Dealt 20 damage to nearby zombies."),

    ("sudo what are all the armor trim patterns?", "armor trim patterns smithing template 1.21",
     "Armor trim patterns (smithing templates):\nCoast, Dune, Eye, Host, Raiser, Rib, Sentry, Shaper, Silence, Snout, Spire, Tide, Vex, Ward, Wayfinder, Wild, Bolt, Flow\n\nMaterials for color: iron, copper, gold, lapis, emerald, diamond, netherite, redstone, amethyst, quartz\n\nApply: Smithing table + template + armor piece + material",
     "https://minecraft.wiki/w/Armor_Trim", ["give {p} minecraft:bolt_armor_trim_smithing_template 1", "give {p} minecraft:flow_armor_trim_smithing_template 1"],
     "Looked up armor trims. Bolt and Flow are the new 1.21 patterns. Gave both templates."),

    ("sudo what biomes have cherry trees?", "cherry grove biome minecraft",
     "Cherry Grove biome:\n- Pink cherry blossom trees\n- Found at Y 100-256 in mountainous areas\n- Pigs, sheep, bees spawn naturally\n- Blocks: cherry_log, cherry_leaves, cherry_planks, cherry_sapling, pink_petals\n- Cherry wood set includes all standard wood items\nLocate: /locate biome minecraft:cherry_grove",
     "https://minecraft.wiki/w/Cherry_Grove", ["locate biome minecraft:cherry_grove"],
     "Looked up cherry grove. Using locate to find the nearest one."),

    ("sudo how does fortune work on different ores?", "fortune enchantment ore drops mechanics",
     "Fortune on ores:\n- Coal, diamond, emerald, lapis, nether quartz, nether gold, copper, amethyst: increases drops\n- Iron, gold ore: Fortune does NOT work (must smelt raw ore)\n- Formula: Each level gives a chance for 1 extra drop\n  - Fortune I: avg 1.33x, Fortune II: avg 1.75x, Fortune III: avg 2.2x\n- Lapis and copper have unique formulas (higher multipliers)\n- Does NOT affect XP drops",
     "https://minecraft.wiki/w/Fortune", ["give {p} minecraft:diamond_pickaxe[enchantments={{fortune:3,efficiency:5,unbreaking:3,mending:1}}] 1"],
     "Looked up fortune mechanics. Works on diamond/emerald/lapis but not iron/gold ore. Gave fortune 3 pickaxe."),
]

# ── Player Info Examples ──────────────────────────────────────────────────

PLAYER_INFO_SCENARIOS = [
    ("sudo build a wall around me", ["fill {x_m10} {y} {z_m10} {x_p10} {y_p4} {z_m10} minecraft:stone", "fill {x_m10} {y} {z_p10} {x_p10} {y_p4} {z_p10} minecraft:stone", "fill {x_m10} {y} {z_m10} {x_m10} {y_p4} {z_p10} minecraft:stone", "fill {x_p10} {y} {z_m10} {x_p10} {y_p4} {z_p10} minecraft:stone"], "Got player position. Built 4 stone walls in a 20x20 square around them."),
    ("sudo teleport me 50 blocks up", ["tp {p} {x} {y_p50} {z}"], "Got player position, teleported 50 blocks up."),
    ("sudo put lava under my feet", ["setblock {x} {y_m1} {z} minecraft:lava"], "Got player position. Placed lava one block below."),
    ("sudo create a beacon at my location", ["setblock {x} {y_m1} {z} minecraft:iron_block", "fill {x_m1} {y_m1} {z_m1} {x_p1} {y_m1} {z_p1} minecraft:iron_block", "setblock {x} {y} {z} minecraft:beacon"], "Got position. Built iron base and placed beacon."),
    ("sudo replace the ground around me with grass", ["fill {x_m5} {y_m1} {z_m5} {x_p5} {y_m1} {z_p5} minecraft:grass_block"], "Got position. Filled a 10x10 area below feet with grass."),
    ("sudo place torches around me in a circle", ["setblock {x_p4} {y} {z} minecraft:torch", "setblock {x_m4} {y} {z} minecraft:torch", "setblock {x} {y} {z_p4} minecraft:torch", "setblock {x} {y} {z_m4} minecraft:torch"], "Got position. Placed 4 torches at cardinal directions."),
    ("sudo tp me to the nearest village center", ["locate structure minecraft:village_plains"], "Got position for reference. Used locate to find nearest village."),
    ("sudo fill the area below me with water", ["fill {x_m3} {y_m5} {z_m3} {x_p3} {y_m1} {z_p3} minecraft:water"], "Got position. Filled 6x5x6 pool below player."),
    ("sudo make a platform of glass under me", ["fill {x_m5} {y_m1} {z_m5} {x_p5} {y_m1} {z_p5} minecraft:glass"], "Got position. Created 10x10 glass platform below feet."),
    ("sudo build an obsidian box around me", ["fill {x_m2} {y} {z_m2} {x_p2} {y_p3} {z_p2} minecraft:obsidian hollow"], "Got position. Built obsidian box around player."),
    ("sudo teleport me to the surface", ["tp {p} {x} 320 {z}"], "Got position. Teleported to build limit so player falls to surface with slow_falling."),
    ("sudo place a campfire at my feet", ["setblock {x} {y} {z} minecraft:campfire"], "Got position. Placed campfire at player location."),
    ("sudo light up this cave", ["setblock {x_p3} {y_p2} {z} minecraft:lantern", "setblock {x_m3} {y_p2} {z} minecraft:lantern", "setblock {x} {y_p2} {z_p3} minecraft:lantern", "setblock {x} {y_p2} {z_m3} minecraft:lantern", "setblock {x} {y_p2} {z} minecraft:lantern"], "Got position. Placed lanterns around the player to illuminate the area."),
    ("sudo surround me with fences", ["fill {x_m3} {y} {z_m3} {x_p3} {y_p1} {z_m3} minecraft:oak_fence replace minecraft:air", "fill {x_m3} {y} {z_p3} {x_p3} {y_p1} {z_p3} minecraft:oak_fence replace minecraft:air"], "Got position. Placed fence perimeter around player."),
    ("sudo clear all blocks 10 blocks above me", ["fill {x_m5} {y_p2} {z_m5} {x_p5} {y_p12} {z_p5} minecraft:air"], "Got position. Cleared 10-block column above player."),
    ("sudo spawn a zombie behind me", ["execute at {p} run summon minecraft:zombie ~ ~ ~-3"], "Got player position. Summoned zombie 3 blocks behind using execute at."),
    ("sudo put a chest with supplies at my location", ["setblock {x} {y} {z} minecraft:chest"], "Got position. Placed chest at player location."),
    ("sudo build stairs going up from where I am", ["setblock {x} {y} {z_p1} minecraft:oak_stairs[facing=south]", "setblock {x} {y_p1} {z_p2} minecraft:oak_stairs[facing=south]", "setblock {x} {y_p2} {z_p3} minecraft:oak_stairs[facing=south]"], "Got position. Built ascending stairs going south."),
    ("sudo make me a fishing spot", ["fill {x_m2} {y_m2} {z_p3} {x_p2} {y_m1} {z_p7} minecraft:water", "give {p} minecraft:fishing_rod 1"], "Got position. Dug pool nearby and gave fishing rod."),
    ("sudo set my spawn point here", ["spawnpoint {p} {x} {y} {z}"], "Got position. Set spawn point at current location."),
]

# ── Server State Examples ─────────────────────────────────────────────────

SERVER_STATE_SCENARIOS = [
    ("sudo give everyone online netherite armor", "Checked online players, giving each full netherite armor."),
    ("sudo announce it's getting dark soon", "Checked time — it's evening. Warning players about night."),
    ("sudo if it's thundering, summon lightning on everyone", "Checked weather. If thundering, striking lightning on all players."),
    ("sudo how many people are online?", "Checked server state. Reporting player count."),
    ("sudo make it night if it's currently day", "Checked time — it was daytime, switching to night."),
    ("sudo give bonus XP to everyone if it's raining", "Checked weather — raining, giving XP bonus."),
    ("sudo tp all players to spawn if it's nighttime", "Checked time — night, teleporting everyone to spawn."),
    ("sudo set weather opposite of current", "Checked weather state and set opposite."),
    ("sudo kill all hostile mobs if there are 3+ players online", "Checked online count. 3+ players → killing hostile mobs."),
    ("sudo announce the time and weather to everyone", "Checked server state. Broadcasting time and weather."),
]

# ── Nearby Entities Examples ──────────────────────────────────────────────

def sim_nearby_entities(types=None):
    all_types = [
        ("zombie", random.randint(1, 5), round(random.uniform(3, 25), 1)),
        ("skeleton", random.randint(1, 3), round(random.uniform(5, 30), 1)),
        ("creeper", random.randint(1, 2), round(random.uniform(4, 20), 1)),
        ("cow", random.randint(2, 8), round(random.uniform(2, 15), 1)),
        ("sheep", random.randint(1, 6), round(random.uniform(3, 20), 1)),
        ("pig", random.randint(1, 4), round(random.uniform(5, 18), 1)),
        ("chicken", random.randint(2, 10), round(random.uniform(1, 12), 1)),
        ("spider", random.randint(1, 3), round(random.uniform(6, 25), 1)),
        ("enderman", 1, round(random.uniform(10, 40), 1)),
        ("villager", random.randint(1, 5), round(random.uniform(3, 20), 1)),
        ("iron_golem", 1, round(random.uniform(5, 15), 1)),
        ("wolf", random.randint(1, 3), round(random.uniform(4, 25), 1)),
    ]
    if types:
        selected = [e for e in all_types if e[0] in types]
    else:
        selected = random.sample(all_types, k=random.randint(2, 5))
    entities = [{"type": f"minecraft:{t}", "count": c, "nearest_distance": d} for t, c, d in selected]
    return {"entities": entities, "total": sum(e["count"] for e in entities)}

NEARBY_SCENARIOS = [
    ("sudo kill all the zombies near me", ["zombie"], ["kill @e[type=minecraft:zombie,distance=..32]"], "Scanned for entities. Found {n} zombies nearby. Killed them all."),
    ("sudo how many mobs are around me?", None, [], "Scanned entities within 32 blocks. Reporting counts."),
    ("sudo clear out the hostile mobs nearby", ["zombie", "skeleton", "creeper", "spider"], ["kill @e[type=minecraft:zombie,distance=..32]", "kill @e[type=minecraft:skeleton,distance=..32]", "kill @e[type=minecraft:creeper,distance=..32]", "kill @e[type=minecraft:spider,distance=..32]"], "Scanned for hostiles. Found and killed all nearby hostile mobs."),
    ("sudo are there any creepers near me?", ["creeper"], [], "Scanned for creepers. Reporting what's nearby."),
    ("sudo kill the closest zombie", ["zombie"], ["kill @e[type=minecraft:zombie,distance=..10,limit=1,sort=nearest]"], "Scanned to confirm zombie presence. Killed the nearest one."),
    ("sudo count the animals around me", ["cow", "sheep", "pig", "chicken"], [], "Scanned for animals. Reporting passive mob counts."),
    ("sudo kill all mobs except villagers in a 50 block radius", None, ["kill @e[type=!minecraft:villager,type=!minecraft:player,distance=..50]"], "Scanned area. Killing everything except villagers and players."),
    ("sudo are there any endermen nearby?", ["enderman"], [], "Scanned for endermen nearby."),
    ("sudo tame any wolves near me", ["wolf"], ["execute as @e[type=minecraft:wolf,distance=..20] run data merge entity @s {Owner:\"{p}\"}"], "Found wolves nearby. Taming them."),
    ("sudo protect me from nearby hostiles", ["zombie", "skeleton", "creeper"], ["kill @e[type=minecraft:zombie,distance=..20]", "kill @e[type=minecraft:skeleton,distance=..20]", "kill @e[type=minecraft:creeper,distance=..20]", "effect give {p} minecraft:resistance 60 2"], "Scanned for hostiles. Killed nearby ones and gave resistance."),
    ("sudo lead me to the nearest villager", ["villager"], [], "Scanned for villagers. Nearest one found — reporting location."),
    ("sudo spawn more of whatever animal is closest", ["cow", "sheep"], ["summon minecraft:cow ~ ~ ~"], "Scanned nearby. Cows closest. Summoned more."),
]

# ── Memory Examples ───────────────────────────────────────────────────────

MEMORY_WRITE_SCENARIOS = [
    ("sudo remember this location as home", "home", "location", {"x": 100, "y": 64, "z": -200}, "Saved current position as 'home'."),
    ("sudo remember my base is here", "base", "location", {"x": -350, "y": 72, "z": 180}, "Saved current position as 'base'."),
    ("sudo remember that I like diamonds", "favorite_item", "preference", "diamonds", "Noted preference for diamonds."),
    ("sudo save this spot as my nether portal", "nether_portal", "location", {"x": 50, "y": 68, "z": 90}, "Saved nether portal location."),
    ("sudo remember I'm building a castle here", "castle", "location", {"x": 200, "y": 80, "z": -100}, "Saved castle build location."),
    ("sudo my favorite color is red", "favorite_color", "preference", "red", "Remembered color preference."),
    ("sudo remember my farm is at 500 70 -300", "farm", "location", {"x": 500, "y": 70, "z": -300}, "Saved farm location from coordinates."),
    ("sudo remember I don't want to be teleported randomly", "no_random_tp", "preference", "no random teleports", "Noted: player prefers no random teleports."),
    ("sudo save this as my mining spot", "mining_spot", "location", {"x": -100, "y": 12, "z": 400}, "Saved mining location."),
    ("sudo remember that Ace13245 is my friend", "friend_ace", "fact", "Ace13245 is a friend", "Saved social fact."),
]

MEMORY_READ_SCENARIOS = [
    ("sudo tp me home", "home", True, "tp {p} {mx} {my} {mz}", "Read memory for 'home'. Teleporting to saved location."),
    ("sudo take me to my base", "base", True, "tp {p} {mx} {my} {mz}", "Read memory for 'base'. Teleporting."),
    ("sudo what do you know about me?", None, True, None, "Read all memories for this player. Listing what I know."),
    ("sudo where's my nether portal?", "nether_portal", True, None, "Read memory for 'nether portal'. Reporting coordinates."),
    ("sudo tp me to my farm", "farm", True, "tp {p} {mx} {my} {mz}", "Read memory for 'farm'. Teleporting."),
    ("sudo do I have a saved home?", "home", False, None, "Checked memories. No 'home' location saved for this player."),
    ("sudo forget my base location", "base", True, None, "Found 'base' memory. Removing it."),
    ("sudo tp me to my mining spot", "mining_spot", True, "tp {p} {mx} {my} {mz}", "Read memory for 'mining_spot'. Teleporting."),
    ("sudo where did I say my castle was?", "castle", True, None, "Read memory for 'castle'. Reporting saved location."),
    ("sudo take me back to where I was building", "castle", True, "tp {p} {mx} {my} {mz}", "Read 'castle' memory. Teleporting to build site."),
]

# ── Chained Tool Call Examples ────────────────────────────────────────────
# These combine multiple tools in one conversation

CHAINED_SCENARIOS = [
    # player_info + nearby_entities + rcon
    ("sudo protect me from everything nearby",
     ["player_info", "nearby_entities", "rcon"],
     "Got position and scanned for hostiles. Killed threats and applied protection."),
    # memory + player_info + rcon
    ("sudo save this spot and build a marker",
     ["player_info", "memory_write", "rcon"],
     "Got position, saved to memory, and built a visible marker."),
    # server_state + rcon (conditional)
    ("sudo if anyone is online, give them all food",
     ["server_state", "rcon"],
     "Checked who's online. Gave food to all players."),
    # memory + rcon (tp to saved location)
    ("sudo tp me home and heal me",
     ["memory_read", "rcon", "rcon"],
     "Read home location from memory. Teleported and healed."),
    # wiki + rcon (look up then execute)
    ("sudo give me whatever the strongest bow setup is",
     ["wiki", "rcon"],
     "Looked up bow enchantments. Gave max enchanted bow."),
    # player_info + server_state + rcon
    ("sudo set up a safe zone around me for the night",
     ["player_info", "server_state", "rcon"],
     "Checked position and time. Built lit shelter for night."),
]


def generate_all():
    examples = []
    idx = 0

    # ── Wiki lookups ──
    for req, query, content, url, cmds, reasoning in WIKI_EXAMPLES:
        player = random.choice(PLAYERS)
        mode = "god" if req.startswith("pray ") else "sudo"
        msgs = [sys_msg(mode), user_msg(f"Player {player}: {req}")]
        msgs.append(tool_call("minecraft.wiki_lookup", {"query": query}))
        msgs.append(tool_result({"content": content, "url": url}))
        for cmd in cmds:
            resolved = cmd.replace("{p}", player)
            msgs.append(tool_call("rcon.execute", {"command": resolved}))
            msgs.append(tool_result({"success": True, "result": rcon_ok(resolved, player)}))
        resp = {"risk_level": 3, "commands": [c.replace("{p}", player) for c in cmds], "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-wiki-{idx:03d}", "wiki_lookup", msgs))
        idx += 1

    # ── Player info ──
    for req, cmds_tmpl, reasoning in PLAYER_INFO_SCENARIOS:
        player = random.choice(PLAYERS)
        pos = rpos()
        mode = "god" if req.startswith("pray ") else "sudo"
        msgs = [sys_msg(mode), user_msg(f"Player {player}: {req}")]
        msgs.append(tool_call("world.player_info", {"player": player}))
        pinfo = {"health": round(random.uniform(10, 20), 1), "position": pos,
                 "inventory_summary": "Diamond sword, iron armor, 32 steak"}
        msgs.append(tool_result(pinfo))

        resolved_cmds = []
        for ct in cmds_tmpl:
            c = ct.replace("{p}", player)
            c = c.replace("{x}", str(pos["x"])).replace("{y}", str(pos["y"])).replace("{z}", str(pos["z"]))
            for delta, tag in [(1, "_p1"), (2, "_p2"), (3, "_p3"), (4, "_p4"), (5, "_p5"), (10, "_p10"), (50, "_p50"),
                               (-1, "_m1"), (-2, "_m2"), (-3, "_m3"), (-5, "_m5"), (-10, "_m10"), (-100, "_minus_100")]:
                c = c.replace(f"{{x{tag}}}", str(pos["x"] + delta))
                c = c.replace(f"{{y{tag}}}", str(pos["y"] + delta))
                c = c.replace(f"{{z{tag}}}", str(pos["z"] + delta))
            resolved_cmds.append(c)

        for cmd in resolved_cmds:
            msgs.append(tool_call("rcon.execute", {"command": cmd}))
            msgs.append(tool_result({"success": True, "result": rcon_ok(cmd, player)}))

        resp = {"risk_level": 3, "commands": resolved_cmds, "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-pinfo-{idx:03d}", "player_info_lookup", msgs))
        idx += 1

    # ── Server state ──
    for req, reasoning in SERVER_STATE_SCENARIOS:
        player = random.choice(PLAYERS)
        mode = "god" if req.startswith("pray ") else "sudo"
        msgs = [sys_msg(mode), user_msg(f"Player {player}: {req}")]
        msgs.append(tool_call("world.server_state", {}))
        state = {
            "time_of_day": random.choice(["morning", "noon", "evening", "night"]),
            "weather": random.choice(["clear", "rain", "thunder"]),
            "online_players": random.sample(PLAYERS, k=random.randint(1, 4)),
            "world_border": 60000000.0,
        }
        msgs.append(tool_result(state))

        # Generate contextual commands based on state
        cmds = []
        if "everyone" in req.lower() or "all" in req.lower():
            for p in state["online_players"]:
                if "netherite" in req: cmds.extend([f"give {p} minecraft:netherite_helmet 1", f"give {p} minecraft:netherite_chestplate 1"])
                elif "food" in req: cmds.append(f"give {p} minecraft:golden_carrot 32")
                elif "XP" in req or "xp" in req: cmds.append(f"xp add {p} 500")
                elif "heal" in req: cmds.append(f"effect give {p} minecraft:instant_health 1 5")
                elif "tp" in req or "spawn" in req: cmds.append(f"tp {p} 0 64 0")
                elif "diamond" in req: cmds.append(f"give {p} minecraft:diamond 1")
        elif "weather" in req.lower():
            opposite = "clear" if state["weather"] in ("rain", "thunder") else "rain"
            cmds.append(f"weather {opposite}")
        elif "night" in req.lower() or "dark" in req.lower():
            cmds.append("time set night" if state["time_of_day"] in ("morning", "noon") else "time set day")
        elif "day" in req.lower():
            cmds.append("time set day" if state["time_of_day"] in ("evening", "night") else "time set night")
        elif "announce" in req.lower() or "how many" in req.lower():
            cmds.append(f"say Server: {len(state['online_players'])} players online, time is {state['time_of_day']}, weather is {state['weather']}")
        elif "hostile" in req.lower() or "kill" in req.lower():
            cmds.append("kill @e[type=!minecraft:player,type=!minecraft:villager,type=!minecraft:iron_golem]")
        elif "lightning" in req.lower():
            for p in state["online_players"]:
                cmds.append(f"execute at {p} run summon minecraft:lightning_bolt")

        if not cmds:
            cmds = ["say Check complete"]

        for cmd in cmds:
            msgs.append(tool_call("rcon.execute", {"command": cmd}))
            msgs.append(tool_result({"success": True, "result": rcon_ok(cmd, player)}))

        resp = {"risk_level": 3, "commands": cmds, "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-sstate-{idx:03d}", "server_state_check", msgs))
        idx += 1

    # ── Nearby entities ──
    for req, entity_filter, cmds_tmpl, reasoning in NEARBY_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {req}")]
        msgs.append(tool_call("world.nearby_entities", {"player": player, "radius": 32}))
        ent_data = sim_nearby_entities(entity_filter)
        msgs.append(tool_result(ent_data))

        cmds = [c.replace("{p}", player) for c in cmds_tmpl]
        n_found = ent_data["total"]
        real_reasoning = reasoning.replace("{n}", str(n_found))

        if cmds:
            for cmd in cmds:
                msgs.append(tool_call("rcon.execute", {"command": cmd}))
                msgs.append(tool_result({"success": True, "result": rcon_ok(cmd, player)}))

        resp = {"risk_level": 3, "commands": cmds, "reasoning": real_reasoning}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-nearby-{idx:03d}", "nearby_entities", msgs))
        idx += 1

    # ── Memory write ──
    for req, key, mtype, value, reasoning in MEMORY_WRITE_SCENARIOS:
        player = random.choice(PLAYERS)
        if mtype == "location" and isinstance(value, dict):
            pos = rpos()
            value = pos
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {req}")]
        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}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-memw-{idx:03d}", "memory_write", msgs))
        idx += 1

    # ── Memory read ──
    for req, key, found, cmd_tmpl, reasoning in MEMORY_READ_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {req}")]

        if key:
            msgs.append(tool_call("memory.read", {"player": player, "key": key}))
        else:
            msgs.append(tool_call("memory.read", {"player": player}))

        if found:
            pos = rpos()
            if key:
                mem = [{"key": key, "type": "location", "value": {"x": pos["x"], "y": pos["y"], "z": pos["z"]}}]
            else:
                mem = [
                    {"key": "home", "type": "location", "value": rpos()},
                    {"key": "favorite_item", "type": "preference", "value": "diamonds"},
                ]
            msgs.append(tool_result({"memories": mem}))

            if cmd_tmpl:
                cmd = cmd_tmpl.replace("{p}", player).replace("{mx}", str(pos["x"])).replace("{my}", str(pos["y"])).replace("{mz}", str(pos["z"]))
                msgs.append(tool_call("rcon.execute", {"command": cmd}))
                msgs.append(tool_result({"success": True, "result": rcon_ok(cmd, player)}))
                resp = {"risk_level": 3, "commands": [cmd], "reasoning": reasoning}
            else:
                resp = {"risk_level": 4, "commands": [], "reasoning": reasoning}
        else:
            msgs.append(tool_result({"memories": []}))
            resp = {"risk_level": 4, "commands": [], "reasoning": reasoning}

        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-memr-{idx:03d}", "memory_read", msgs))
        idx += 1

    # ── Chained tool calls ──
    for req, tools, reasoning in CHAINED_SCENARIOS:
        player = random.choice(PLAYERS)
        pos = rpos()
        msgs = [sys_msg("sudo"), user_msg(f"Player {player}: {req}")]

        for t in tools:
            if t == "player_info":
                msgs.append(tool_call("world.player_info", {"player": player}))
                msgs.append(tool_result({"health": 18.5, "position": pos, "inventory_summary": "Diamond gear"}))
            elif t == "nearby_entities":
                msgs.append(tool_call("world.nearby_entities", {"player": player, "radius": 32}))
                msgs.append(tool_result(sim_nearby_entities(["zombie", "skeleton"])))
            elif t == "server_state":
                msgs.append(tool_call("world.server_state", {}))
                msgs.append(tool_result({"time_of_day": "night", "weather": "clear", "online_players": [player, "Ace13245"], "world_border": 60000000.0}))
            elif t == "memory_write":
                msgs.append(tool_call("memory.write", {"player": player, "type": "location", "key": "marker", "value": pos}))
                msgs.append(tool_result({"ok": True, "key": "marker"}))
            elif t == "memory_read":
                msgs.append(tool_call("memory.read", {"player": player, "key": "home"}))
                home = rpos()
                msgs.append(tool_result({"memories": [{"key": "home", "type": "location", "value": home}]}))
                pos = home  # use home coords for subsequent commands
            elif t == "wiki":
                msgs.append(tool_call("minecraft.wiki_lookup", {"query": "best bow enchantments 1.21"}))
                msgs.append(tool_result({"content": "Bow: Power V, Infinity or Mending, Unbreaking III, Flame I, Punch II. Infinity and Mending are mutually exclusive.", "url": "https://minecraft.wiki/w/Bow"}))
            elif t == "rcon":
                cmd = f"give {player} minecraft:diamond 1"
                if "heal" in req: cmd = f"effect give {player} minecraft:instant_health 1 5"
                elif "protect" in req: cmd = f"kill @e[type=minecraft:zombie,distance=..32]"
                elif "marker" in req or "build" in req: cmd = f"setblock {pos['x']} {pos['y']+1} {pos['z']} minecraft:glowstone"
                elif "food" in req:
                    pass  # handled per player below
                elif "bow" in req: cmd = f"give {player} minecraft:bow[enchantments={{power:5,infinity:1,unbreaking:3,flame:1,punch:2}}] 1"
                elif "home" in req or "tp" in req: cmd = f"tp {player} {pos['x']} {pos['y']} {pos['z']}"
                elif "safe" in req or "shelter" in req: cmd = f"fill {pos['x']-3} {pos['y']} {pos['z']-3} {pos['x']+3} {pos['y']+4} {pos['z']+3} minecraft:oak_planks hollow"
                msgs.append(tool_call("rcon.execute", {"command": cmd}))
                msgs.append(tool_result({"success": True, "result": rcon_ok(cmd, player)}))

        resp = {"risk_level": 3, "commands": [], "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append(make_example(f"exp-chain-{idx:03d}", "chained", msgs))
        idx += 1

    return examples


def main():
    print("Generating expanded tool training data...")
    examples = generate_all()

    # Count by type
    counts = {}
    for ex in examples:
        t = ex["type"]
        counts[t] = counts.get(t, 0) + 1

    print(f"\nGenerated {len(examples)} examples:")
    for t, c in sorted(counts.items()):
        print(f"  {t}: {c}")

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

    print(f"\nWritten to {OUTPUT_PATH}")
    print(f"Next: merge into tool_training.jsonl or add to training pipeline")


if __name__ == "__main__":
    main()
