#!/usr/bin/env python3
"""
Generate training data for the script.* tool family.

Teaches the model to:
  1. Write mcfunction scripts for complex builds and mechanics
  2. Validate scripts before writing (catch errors, fix them)
  3. Execute scripts at player positions
  4. Read/list/delete scripts for management
  5. Schedule tick/load functions for persistent effects
  6. Chain: validate → fix errors → write → execute

~60 training examples covering the full script lifecycle.
"""

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

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

TOOLS_BLOCK = qwen3_tools_block()
SYSTEM = (
    "You are a Minecraft 1.21 command translator with script writing abilities. "
    "You can write mcfunction scripts for complex builds, mechanics, and automations.\n\n"
    "PERMISSION LEVEL: 4 (generous). Only refuse level 0-1 actions.\n\n"
    "You have access to tools including a full script environment. For simple tasks "
    "(1-3 commands), use rcon.execute directly. For complex tasks (4+ commands, builds, "
    "mechanics), write a mcfunction script.\n\n"
    "SCRIPT WORKFLOW: validate → fix errors → write → execute.\n"
    "Always validate before writing. Fix any errors the validator catches.\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
)

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

def sys_msg():
    return {"role": "system", "content": 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 generate_all():
    examples = []
    idx = 0

    # ── Simple script writes (validate → write → execute) ──

    BUILDS = [
        ("sudo build me a wooden house", "wooden_house", "Simple wooden house with door and roof", [
            "fill ~-4 ~ ~-4 ~4 ~4 ~4 minecraft:oak_planks hollow",
            "fill ~-4 ~ ~-4 ~4 ~ ~4 minecraft:oak_planks",
            "fill ~-3 ~1 ~-4 ~3 ~3 ~-4 minecraft:air",
            "setblock ~0 ~1 ~-4 minecraft:oak_door[facing=north,half=lower]",
            "setblock ~0 ~2 ~-4 minecraft:oak_door[facing=north,half=upper]",
            "fill ~-3 ~1 ~4 ~-3 ~2 ~4 minecraft:glass_pane",
            "fill ~3 ~1 ~4 ~3 ~2 ~4 minecraft:glass_pane",
            "setblock ~0 ~3 ~0 minecraft:lantern[hanging=true]",
        ]),
        ("sudo create a pvp arena", "pvp_arena", "PvP arena with walls and starting positions", [
            "fill ~-15 ~-1 ~-15 ~15 ~-1 ~15 minecraft:smooth_stone",
            "fill ~-15 ~ ~-15 ~15 ~5 ~-15 minecraft:stone_bricks",
            "fill ~-15 ~ ~15 ~15 ~5 ~15 minecraft:stone_bricks",
            "fill ~-15 ~ ~-15 ~-15 ~5 ~15 minecraft:stone_bricks",
            "fill ~15 ~ ~-15 ~15 ~5 ~15 minecraft:stone_bricks",
            "fill ~-14 ~ ~-14 ~14 ~4 ~14 minecraft:air",
            "setblock ~-10 ~ ~0 minecraft:red_concrete",
            "setblock ~10 ~ ~0 minecraft:blue_concrete",
            "fill ~-1 ~-1 ~-1 ~1 ~-1 ~1 minecraft:glowstone",
        ]),
        ("sudo make a nether portal room", "portal_room", "Decorated nether portal room", [
            "fill ~-5 ~ ~-5 ~5 ~6 ~5 minecraft:blackstone hollow",
            "fill ~-4 ~ ~-4 ~4 ~5 ~4 minecraft:air",
            "fill ~-1 ~1 ~0 ~1 ~5 ~0 minecraft:obsidian",
            "fill ~0 ~1 ~0 ~0 ~4 ~0 minecraft:air",
            "setblock ~0 ~1 ~0 minecraft:nether_portal[axis=x]",
            "setblock ~0 ~2 ~0 minecraft:nether_portal[axis=x]",
            "setblock ~0 ~3 ~0 minecraft:nether_portal[axis=x]",
            "fill ~-4 ~-1 ~-4 ~4 ~-1 ~4 minecraft:polished_blackstone_bricks",
            "setblock ~-3 ~1 ~-3 minecraft:soul_lantern",
            "setblock ~3 ~1 ~-3 minecraft:soul_lantern",
            "setblock ~-3 ~1 ~3 minecraft:soul_lantern",
            "setblock ~3 ~1 ~3 minecraft:soul_lantern",
        ]),
        ("sudo build a watchtower", "watchtower", "Tall watchtower with ladder access", [
            "fill ~-2 ~ ~-2 ~2 ~10 ~2 minecraft:cobblestone hollow",
            "fill ~-1 ~ ~-1 ~1 ~9 ~1 minecraft:air",
            "fill ~-2 ~10 ~-2 ~2 ~10 ~2 minecraft:oak_planks",
            "fill ~-2 ~11 ~-2 ~2 ~11 ~2 minecraft:oak_fence",
            "fill ~-1 ~11 ~-1 ~1 ~11 ~1 minecraft:air",
            "fill ~2 ~1 ~0 ~2 ~9 ~0 minecraft:ladder[facing=west]",
            "setblock ~0 ~12 ~0 minecraft:lantern",
        ]),
        ("sudo build a farm plot with water", "farm_plot", "9x9 farm with central water and tilled soil", [
            "fill ~-4 ~-1 ~-4 ~4 ~-1 ~4 minecraft:farmland",
            "setblock ~0 ~-1 ~0 minecraft:water",
            "fill ~-4 ~ ~-4 ~4 ~ ~4 minecraft:air",
            "fill ~-5 ~-1 ~-5 ~5 ~-1 ~-5 minecraft:oak_fence",
            "fill ~-5 ~-1 ~5 ~5 ~-1 ~5 minecraft:oak_fence",
            "fill ~-5 ~-1 ~-5 ~-5 ~-1 ~5 minecraft:oak_fence",
            "fill ~5 ~-1 ~-5 ~5 ~-1 ~5 minecraft:oak_fence",
            "setblock ~-5 ~-1 ~0 minecraft:oak_fence_gate[facing=east]",
        ]),
        ("sudo make a mob grinder platform", "mob_grinder", "Mob spawning platform with water channels", [
            "fill ~-8 ~ ~-8 ~8 ~ ~8 minecraft:cobblestone",
            "fill ~-7 ~1 ~-7 ~7 ~3 ~7 minecraft:air",
            "fill ~-8 ~1 ~-8 ~8 ~1 ~-8 minecraft:cobblestone",
            "fill ~-8 ~1 ~8 ~8 ~1 ~8 minecraft:cobblestone",
            "fill ~-8 ~1 ~-8 ~-8 ~1 ~8 minecraft:cobblestone",
            "fill ~8 ~1 ~-8 ~8 ~1 ~8 minecraft:cobblestone",
            "setblock ~0 ~-1 ~0 minecraft:air",
            "setblock ~0 ~-2 ~0 minecraft:hopper",
            "fill ~-8 ~ ~0 ~-1 ~ ~0 minecraft:water",
            "fill ~1 ~ ~0 ~8 ~ ~0 minecraft:water",
        ]),
        ("sudo create a trophy room", "trophy_room", "Display room with item frames and lighting", [
            "fill ~-6 ~ ~-6 ~6 ~4 ~6 minecraft:smooth_quartz hollow",
            "fill ~-5 ~ ~-5 ~5 ~3 ~5 minecraft:air",
            "fill ~-6 ~-1 ~-6 ~6 ~-1 ~6 minecraft:polished_deepslate",
            "fill ~-5 ~2 ~-6 ~5 ~2 ~-6 minecraft:item_frame[facing=south]",
            "fill ~-5 ~2 ~6 ~5 ~2 ~6 minecraft:item_frame[facing=north]",
            "setblock ~-4 ~3 ~0 minecraft:lantern[hanging=true]",
            "setblock ~0 ~3 ~0 minecraft:lantern[hanging=true]",
            "setblock ~4 ~3 ~0 minecraft:lantern[hanging=true]",
        ]),
    ]

    for req, name, desc, cmds in BUILDS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]

        # Validate first
        msgs.append(tool_call("script.validate", {"commands": cmds}))
        msgs.append(tool_result({"valid": True, "total": len(cmds), "passed": len(cmds), "errors": []}))

        # Write
        msgs.append(tool_call("script.write", {"name": name, "commands": cmds, "description": desc}))
        msgs.append(tool_result({"ok": True, "path": f"mortdecai:{name}", "lines": len(cmds)}))

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

        resp = {"risk_level": 3, "commands": [f"function mortdecai:{name}"], "reasoning": f"Wrote and executed '{name}' script ({len(cmds)} commands) at {player}'s position."}
        msgs.append(final_response(resp))
        examples.append({"id": f"script-build-{idx:03d}", "source": "script_training", "type": "script_write_execute", "messages": msgs})
        idx += 1

    # ── Validate-fix-write cycle (model catches and fixes errors) ──

    ERROR_SCENARIOS = [
        ("sudo build a glass bridge", "glass_bridge", [
            "fill ~0 ~-1 ~0 ~0 ~-1 ~20 minecraft:glass",
            "fill ~-1 ~ ~0 ~-1 ~ ~20 minecraft:glass_pane",
            "fill ~1 ~ ~0 ~1 ~ ~20 minecraft:glass_paine",  # typo
        ], 2, "glass_paine", "Unknown block type: minecraft:glass_paine", [
            "fill ~0 ~-1 ~0 ~0 ~-1 ~20 minecraft:glass",
            "fill ~-1 ~ ~0 ~-1 ~ ~20 minecraft:glass_pane",
            "fill ~1 ~ ~0 ~1 ~ ~20 minecraft:glass_pane",
        ]),
        ("sudo create an enchanting setup", "enchanting_room", [
            "setblock ~0 ~ ~0 minecraft:enchanting_table",
            "fill ~-2 ~ ~-2 ~2 ~1 ~2 minecraft:bookshelf",
            "fill ~-1 ~ ~-1 ~1 ~1 ~1 minecraft:air",
            "setblock ~0 ~2 ~0 minecraft:lantern",
            "effect give @s minecraft:luck 600",  # missing amplifier — not an error but let's say
        ], 4, "minecraft:luck", "Unknown effect: minecraft:luck", [
            "setblock ~0 ~ ~0 minecraft:enchanting_table",
            "fill ~-2 ~ ~-2 ~2 ~1 ~2 minecraft:bookshelf",
            "fill ~-1 ~ ~-1 ~1 ~1 ~1 minecraft:air",
            "setblock ~0 ~2 ~0 minecraft:lantern",
        ]),
        ("sudo make a redstone door", "redstone_door", [
            "setblock ~2 ~1 ~0 minecraft:sticky_pistons[facing=west]",  # wrong ID
            "setblock ~2 ~ ~0 minecraft:redstone_block",
        ], 0, "minecraft:sticky_pistons", "Unknown block: minecraft:sticky_pistons", [
            "setblock ~2 ~1 ~0 minecraft:sticky_piston[facing=west]",
            "setblock ~2 ~ ~0 minecraft:redstone_block",
        ]),
    ]

    for req, name, bad_cmds, err_line, err_cmd_fragment, err_msg, fixed_cmds in ERROR_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]

        # First validation fails
        msgs.append(tool_call("script.validate", {"commands": bad_cmds}))
        msgs.append(tool_result({
            "valid": False, "total": len(bad_cmds), "passed": len(bad_cmds) - 1,
            "errors": [{"line": err_line + 1, "command": bad_cmds[err_line], "error": err_msg}]
        }))

        # Fix and re-validate
        msgs.append(tool_call("script.validate", {"commands": fixed_cmds}))
        msgs.append(tool_result({"valid": True, "total": len(fixed_cmds), "passed": len(fixed_cmds), "errors": []}))

        # Write fixed version
        msgs.append(tool_call("script.write", {"name": name, "commands": fixed_cmds}))
        msgs.append(tool_result({"ok": True, "path": f"mortdecai:{name}", "lines": len(fixed_cmds)}))

        # Execute
        msgs.append(tool_call("script.execute", {"name": name, "as_player": player}))
        msgs.append(tool_result({"ok": True, "result": "Executed commands"}))

        resp = {"risk_level": 3, "commands": [f"function mortdecai:{name}"], "reasoning": f"Validation caught an error. Fixed and wrote '{name}' script."}
        msgs.append(final_response(resp))
        examples.append({"id": f"script-fix-{idx:03d}", "source": "script_training", "type": "script_validate_fix", "messages": msgs})
        idx += 1

    # ── Script management (list, read, delete) ──

    MGMT_SCENARIOS = [
        ("sudo show me all my scripts", "list", None),
        ("sudo what's in the pvp_arena script?", "read", "pvp_arena"),
        ("sudo delete the old test script", "delete", "test_script"),
        ("sudo what scripts do I have?", "list", None),
        ("sudo show me the watchtower script", "read", "watchtower"),
        ("sudo remove the mob_grinder script", "delete", "mob_grinder"),
    ]

    for req, action, target in MGMT_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]

        if action == "list":
            msgs.append(tool_call("script.list", {}))
            msgs.append(tool_result({"scripts": [
                {"name": "wooden_house", "lines": 8, "scheduled": "none"},
                {"name": "pvp_arena", "lines": 9, "scheduled": "none"},
                {"name": "portal_room", "lines": 12, "scheduled": "none"},
                {"name": "ambient_particles", "lines": 3, "scheduled": "tick"},
            ]}))
            resp = {"risk_level": 4, "commands": [], "reasoning": "Listed all scripts in the datapack."}
        elif action == "read":
            msgs.append(tool_call("script.read", {"name": target}))
            msgs.append(tool_result({"ok": True, "commands": ["fill ~-15 ~-1 ~-15 ~15 ~-1 ~15 minecraft:smooth_stone", "# ... more commands"], "lines": 9}))
            resp = {"risk_level": 4, "commands": [], "reasoning": f"Read the '{target}' script contents."}
        elif action == "delete":
            msgs.append(tool_call("script.delete", {"name": target}))
            msgs.append(tool_result({"ok": True}))
            resp = {"risk_level": 3, "commands": [], "reasoning": f"Deleted the '{target}' script."}

        msgs.append(final_response(resp))
        examples.append({"id": f"script-mgmt-{idx:03d}", "source": "script_training", "type": "script_management", "messages": msgs})
        idx += 1

    # ── Tick/load scheduling ──

    SCHEDULE_SCENARIOS = [
        ("sudo make ambient particles around spawn forever", "ambient_particles", "tick", [
            "execute at @a run particle minecraft:cherry_leaves ~ ~2 ~ 3 1 3 0.02 5",
        ]),
        ("sudo set up scoreboards when the server starts", "init_scoreboards", "load", [
            "scoreboard objectives add deaths deathCount Deaths",
            "scoreboard objectives add kills playerKillCount Kills",
            "scoreboard objectives setdisplay sidebar kills",
        ]),
        ("sudo make it rain XP orbs on everyone every tick", "xp_rain", "tick", [
            "execute at @a run summon minecraft:experience_orb ~ ~3 ~ {Value:1}",
        ]),
        ("sudo initialize gamerules on server load", "init_gamerules", "load", [
            "gamerule keepInventory true",
            "gamerule doDaylightCycle true",
            "gamerule announceAdvancements true",
            "gamerule playersSleepingPercentage 50",
        ]),
    ]

    for req, name, sched_type, cmds in SCHEDULE_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]

        # Validate
        msgs.append(tool_call("script.validate", {"commands": cmds}))
        msgs.append(tool_result({"valid": True, "total": len(cmds), "passed": len(cmds), "errors": []}))

        # Write
        msgs.append(tool_call("script.write", {"name": name, "commands": cmds, "description": req}))
        msgs.append(tool_result({"ok": True, "path": f"mortdecai:{name}", "lines": len(cmds)}))

        # Schedule
        msgs.append(tool_call("script.schedule", {"name": name, "type": sched_type}))
        msgs.append(tool_result({"ok": True}))

        resp = {"risk_level": 3 if sched_type == "load" else 2, "commands": [f"function mortdecai:{name}"], "reasoning": f"Wrote '{name}' and scheduled it to run on {sched_type}."}
        msgs.append(final_response(resp))
        examples.append({"id": f"script-sched-{idx:03d}", "source": "script_training", "type": "script_schedule", "messages": msgs})
        idx += 1

    # ── Direct rcon for simple tasks (model should NOT use scripts for 1-2 commands) ──

    SIMPLE = [
        ("sudo give me a diamond", ["give {p} minecraft:diamond 1"], "Simple give — no script needed."),
        ("sudo set time to day", ["time set day"], "Single command, no script needed."),
        ("sudo kill nearby zombies", ["kill @e[type=minecraft:zombie,distance=..30]"], "One command, direct execution is better."),
    ]

    for req, cmds, reasoning in SIMPLE:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]
        for cmd in cmds:
            resolved = cmd.replace("{p}", player)
            msgs.append(tool_call("rcon.execute", {"command": resolved}))
            msgs.append(tool_result({"success": True, "result": "Command executed successfully"}))
        resp = {"risk_level": 3, "commands": [c.replace("{p}", player) for c in cmds], "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append({"id": f"script-simple-{idx:03d}", "source": "script_training", "type": "direct_rcon", "messages": msgs})
        idx += 1

    # ── Edit cycle: read → modify → validate → write ──

    EDIT_SCENARIOS = [
        ("sudo add lighting to the pvp arena", "pvp_arena",
         ["fill ~-15 ~-1 ~-15 ~15 ~-1 ~15 minecraft:smooth_stone", "fill ~-15 ~ ~-15 ~15 ~5 ~-15 minecraft:stone_bricks"],
         ["fill ~-15 ~-1 ~-15 ~15 ~-1 ~15 minecraft:smooth_stone", "fill ~-15 ~ ~-15 ~15 ~5 ~-15 minecraft:stone_bricks",
          "setblock ~-10 ~4 ~0 minecraft:lantern[hanging=true]", "setblock ~0 ~4 ~0 minecraft:lantern[hanging=true]", "setblock ~10 ~4 ~0 minecraft:lantern[hanging=true]"],
         "Read existing arena script, added hanging lanterns for lighting."),
        ("sudo add a roof to the wooden house", "wooden_house",
         ["fill ~-4 ~ ~-4 ~4 ~4 ~4 minecraft:oak_planks hollow"],
         ["fill ~-4 ~ ~-4 ~4 ~4 ~4 minecraft:oak_planks hollow",
          "fill ~-5 ~4 ~-5 ~5 ~4 ~5 minecraft:oak_slab", "fill ~-4 ~5 ~-4 ~4 ~5 ~4 minecraft:oak_slab"],
         "Read existing house script, added oak slab roof."),
    ]

    for req, name, old_cmds, new_cmds, reasoning in EDIT_SCENARIOS:
        player = random.choice(PLAYERS)
        msgs = [sys_msg(), user_msg(f"Player {player}: {req}")]

        # Read existing
        msgs.append(tool_call("script.read", {"name": name}))
        msgs.append(tool_result({"ok": True, "commands": old_cmds, "lines": len(old_cmds)}))

        # Validate new version
        msgs.append(tool_call("script.validate", {"commands": new_cmds}))
        msgs.append(tool_result({"valid": True, "total": len(new_cmds), "passed": len(new_cmds), "errors": []}))

        # Overwrite
        msgs.append(tool_call("script.write", {"name": name, "commands": new_cmds}))
        msgs.append(tool_result({"ok": True, "path": f"mortdecai:{name}", "lines": len(new_cmds)}))

        # Execute
        msgs.append(tool_call("script.execute", {"name": name, "as_player": player}))
        msgs.append(tool_result({"ok": True, "result": "Executed commands"}))

        resp = {"risk_level": 3, "commands": [f"function mortdecai:{name}"], "reasoning": reasoning}
        msgs.append(final_response(resp))
        examples.append({"id": f"script-edit-{idx:03d}", "source": "script_training", "type": "script_edit", "messages": msgs})
        idx += 1

    return examples


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

    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}")

    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}")


if __name__ == "__main__":
    main()