scripts/generate_code_mutation_sft_seed.py

from __future__ import annotations

import argparse
import ast
import json
import re
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from life_game.code_extensions import extension_briefs_for_mode  # noqa: E402
from life_game.code_mutation import parse_method_replacement_spec  # noqa: E402
from life_game.game import GAME_REGISTRY, SANDBOX_MODE  # noqa: E402

SCHEMA = "signal-garden-code-mutation-sft/v1"
DEFAULT_OUTPUT = "annotations/code_mutation_sft_seed.jsonl"
SYSTEM_PROMPT = (
    "You are an expert Python game engineer. Return only JSON with complete Python method replacements. "
    "Build material game mechanics as deterministic local changes: store new state in game.data or existing "
    "GameModel masks, update player-visible messages/progress, prefer complete new()/command()/advance()/progress() "
    "method replacements, and avoid app UI, renderer, replay, dependency, or multi-file architecture work unless "
    "the prompt explicitly requests it."
)
RESPONSE_CONTRACT = """Return JSON only with this schema:
{
  "summary": "short imperative summary",
  "rationale": "why the code change helps the game",
  "replacements": [
    {
      "path": "life_game/games/example.py",
      "class_name": "ExampleGame",
      "method_name": "advance",
      "content": "    def advance(self, model, game, rule):\\n        ..."
    }
  ],
  "tests": ["uv run pytest tests/test_game.py"]
}

For method replacements:
- content must be the complete replacement method, including the def line.
- Replace the required methods and keep JSON valid.
- Do not invent helper methods or new imports unless the returned replacements define everything needed.
- Store new mechanic state in game.data.
- Directly update existing GameModel fields such as bullets, enemies, towers, or pickups when a mechanic needs new entities.
- Keep behavior deterministic from ticks, score, health, commands, or existing masks.
- Include tests that would verify the changed mechanic."""


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate seed ShareGPT SFT rows for Signal Garden code mutations.")
    parser.add_argument("--output", default=DEFAULT_OUTPUT)
    parser.add_argument("--count", type=int, default=100)
    args = parser.parse_args()

    rows = build_dataset(max(1, int(args.count)))
    output = ROOT / args.output
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")
    try:
        display = output.relative_to(ROOT)
    except ValueError:
        display = output
    print(f"Wrote {len(rows)} rows to {display}")


def build_dataset(count: int) -> list[dict[str, Any]]:
    briefs = _brief_inventory()
    rows: list[dict[str, Any]] = []
    for index in range(count):
        mode, brief = briefs[index % len(briefs)]
        variant = index // len(briefs)
        repair = index % 5 == 4
        assistant = _assistant_target(mode, brief, variant)
        _validate_assistant_target(assistant)
        user = _user_prompt(mode, brief, variant, repair)
        rows.append(
            {
                "schema": SCHEMA,
                "id": f"code-mutation-sft-{index + 1:03d}",
                "split": "eval" if index % 10 == 9 else "train",
                "source": "synthetic_from_signal_garden_extension_briefs",
                "metadata": {
                    "mode": mode,
                    "extension_id": brief.extension_id,
                    "title": brief.title,
                    "required_methods": list(brief.required_methods or ("command", "advance")),
                    "target_change_lines": brief.target_change_lines,
                    "repair_style_prompt": repair,
                    "curation_note": (
                        "Seed row for LoRA/SFT format and mechanic alignment. Review generated code targets before "
                        "treating this as production patch ground truth."
                    ),
                },
                "conversations": [
                    {"from": "system", "value": SYSTEM_PROMPT},
                    {"from": "human", "value": user},
                    {"from": "gpt", "value": json.dumps(assistant, sort_keys=True)},
                ],
            }
        )
    return rows


def _brief_inventory() -> list[tuple[str, Any]]:
    inventory: list[tuple[str, Any]] = []
    for mode in GAME_REGISTRY:
        if mode == SANDBOX_MODE:
            continue
        for brief in extension_briefs_for_mode(mode):
            inventory.append((mode, brief))
    if not inventory:
        raise RuntimeError("No extension briefs available.")
    return inventory


def _user_prompt(mode: str, brief: Any, variant: int, repair: bool) -> str:
    game = GAME_REGISTRY[mode]
    class_name = game.__class__.__name__
    path = _mode_path(game.__class__.__module__)
    methods = ", ".join(brief.required_methods or ("command", "advance"))
    notes = "\n".join(f"    - {note}" for note in brief.implementation_notes)
    tests = "\n".join(f"    - {test}" for test in brief.suggested_tests)
    risks = "\n".join(f"    - {risk}" for risk in brief.risk_notes)
    source = _numbered_training_source(class_name, mode, path)
    repair_context = ""
    if repair:
        repair_context = """

Previous output failed validation:
- It copied existing methods without changing behavior.
- It missed required semantic terms from the extension brief.
- It returned prose before JSON.

Return the corrected JSON only."""
    return f"""Mode: {mode}
Objective: {game.objective}
Developer intent: {brief.player_facing_goal}
Extension id: {brief.extension_id}
Variant index: {variant}

{RESPONSE_CONTRACT}

Required replacement methods: {methods}

Implementation notes:
{notes}

Suggested tests:
{tests}

Risk notes:
{risks}

Numbered {path}:
```python
{source}
```{repair_context}""".strip()


def _assistant_target(mode: str, brief: Any, variant: int) -> dict[str, Any]:
    game = GAME_REGISTRY[mode]
    class_name = game.__class__.__name__
    path = _mode_path(game.__class__.__module__)
    methods = tuple(brief.required_methods or ("command", "advance"))
    state_key = _state_key(brief.extension_id)
    replacements = [
        {
            "path": path,
            "class_name": class_name,
            "method_name": method,
            "content": _method_content(method, mode, brief, state_key, variant),
        }
        for method in methods
    ]
    return {
        "summary": f"Add {brief.title.lower()}",
        "rationale": brief.player_facing_goal,
        "replacements": replacements,
        "tests": [
            "uv run pytest tests/test_game.py",
            "uv run pytest tests/test_code_mutation.py",
        ],
    }


def _method_content(method: str, mode: str, brief: Any, state_key: str, variant: int) -> str:
    label = _message_label(brief.title)
    charge_max = 3 + (variant % 3)
    cooldown = 5 + (variant % 4)
    if method == "new":
        return f"""    def new(self, size: int, rng: np.random.Generator) -> GameModel:
        enemies = np.zeros((size, size), dtype=bool)
        player = (size // 2, size // 2)
        data = {{
            "target": max(3, size // 4),
            "{state_key}_charges": 1,
            "{state_key}_cooldown": 0,
            "{state_key}_ticks": 0,
            "{state_key}_mode": "{_primary_word(brief)}",
        }}
        return GameModel(
            mode=self.mode,
            enemies=enemies,
            player=player,
            health=6,
            max_health=6,
            data=data,
            message="{label} READY",
        )
"""
    if method == "command":
        return f"""    def command(self, model: BoardModel, game: GameModel, command: str) -> tuple[BoardModel, GameModel]:
        data = dict(game.data)
        if command in DIRECTIONS:
            return model, move_player(game, command, model.grid.shape, wrap=True)
        if command == "fire":
            cooldown = int(data.get("{state_key}_cooldown", 0))
            charges = int(data.get("{state_key}_charges", 0))
            modes = ("blocker", "zapper", "slow", "shield", "overcharge", "safe_zone", "heavy", "phase")
            current = str(data.get("{state_key}_mode", modes[0]))
            if cooldown > 0 or charges <= 0:
                data["{state_key}_mode"] = current
                return model, replace(game, data=data, message="{label} COOLING")
            next_index = (modes.index(current) + 1) % len(modes) if current in modes else 0
            data["{state_key}_mode"] = modes[next_index]
            data["{state_key}_ticks"] = {charge_max}
            data["{state_key}_cooldown"] = {cooldown}
            data["{state_key}_charges"] = max(0, charges - 1)
            return model, emit_sound(replace(game, data=data, message="{label} {str(brief.extension_id).upper()}"), "zap")
        return model, game
"""
    if method == "advance":
        terms = _semantic_terms(brief)
        comment = ", ".join(terms[:5])
        return f"""    def advance(self, model: BoardModel, game: GameModel, rule: LifeRule) -> tuple[BoardModel, GameModel]:
        data = dict(game.data)
        enemies = game.enemies.copy() if game.enemies is not None else np.zeros(model.grid.shape, dtype=bool)
        health = game.health
        score = game.score
        events: list[str] = []

        # Deterministic extension state: {comment}.
        active_ticks = max(0, int(data.get("{state_key}_ticks", 0)) - 1)
        data["{state_key}_ticks"] = active_ticks
        data["{state_key}_cooldown"] = max(0, int(data.get("{state_key}_cooldown", 0)) - 1)
        mode = str(data.get("{state_key}_mode", "{_primary_word(brief)}"))

        if active_ticks > 0 and game.player is not None:
            row, col = game.player
            radius = 1 + int(mode in {{"zapper", "overcharge", "safe_zone"}})
            rr, cc = np.ogrid[: enemies.shape[0], : enemies.shape[1]]
            mask = (rr - row) ** 2 + (cc - col) ** 2 <= radius * radius
            if mode in {{"zapper", "overcharge", "safe_zone"}}:
                cleared = int(enemies[mask].sum())
                enemies[mask] = False
                score += cleared * 5
                if cleared:
                    events.append("zap")
            elif mode in {{"shield", "blocker"}} and enemies[game.player]:
                enemies[game.player] = False
                events.append("deploy")
            elif mode in {{"slow", "phase", "heavy"}} and game.ticks % 2 == 0:
                data["{state_key}_cooldown"] = max(data["{state_key}_cooldown"], 1)

        if game.player is not None and enemies[game.player]:
            enemies[game.player] = False
            health = max(0, health - 1)
            events.append("damage")

        next_game = replace(game, enemies=enemies, health=health, score=score, data=data, message="{label} ACTIVE")
        return model, emit_sound(next_game, *events) if events else next_game
"""
    if method == "progress":
        return f"""    def progress(self, game: GameModel) -> tuple[str, str, float | None]:
        target = int(game.data.get("target", 6))
        active = int(game.data.get("{state_key}_ticks", 0))
        progress = min(target, game.score)
        label = "{label}"
        return label, f"{{progress}}/{{target}} + active {{active}}", min(1.0, progress / max(1, target))
"""
    if method == "is_complete":
        return """    def is_complete(self, game: GameModel) -> bool:
        return game.score >= int(game.data.get("target", 6))
"""
    return f"""    def {method}(self, *args, **kwargs):
        raise NotImplementedError("Training seed does not define this method.")
"""


def _validate_assistant_target(payload: dict[str, Any]) -> None:
    parse_method_replacement_spec(json.dumps(payload))
    for replacement in payload["replacements"]:
        method_name = replacement["method_name"]
        content = replacement["content"]
        first = content.strip().splitlines()[0]
        if not first.lstrip().startswith(f"def {method_name}("):
            raise RuntimeError(f"Bad method content for {method_name}.")
        ast.parse("class _TrainingTarget:\n" + _indent_method(content))


def _indent_method(content: str) -> str:
    lines = content.strip("\n").splitlines()
    if lines and lines[0].startswith("def "):
        lines = [f"    {line}" if line else "" for line in lines]
    return "\n".join(lines) + "\n"


def _numbered_training_source(class_name: str, mode: str, path: str) -> str:
    source = f'''from __future__ import annotations
from dataclasses import replace
import numpy as np
from life_game.audio import emit_sound
from life_game.board import BoardModel
from life_game.rules import LifeRule
from .base import CLICK_NONE, GameModel, GameMode
from .common import DIRECTIONS, move_player


class {class_name}(GameMode):
    mode = "{mode}"
    click_policy = CLICK_NONE

    def new(self, size: int, rng: np.random.Generator) -> GameModel:
        enemies = np.zeros((size, size), dtype=bool)
        return GameModel(mode=self.mode, enemies=enemies, player=(size // 2, size // 2), health=6, max_health=6, data={{"target": 6}}, message="READY")

    def command(self, model: BoardModel, game: GameModel, command: str) -> tuple[BoardModel, GameModel]:
        if command in DIRECTIONS:
            return model, move_player(game, command, model.grid.shape, wrap=True)
        return model, game

    def advance(self, model: BoardModel, game: GameModel, rule: LifeRule) -> tuple[BoardModel, GameModel]:
        return model, game
'''
    return "\n".join(f"{index:04d}: {line}" for index, line in enumerate(source.splitlines(), start=1))


def _mode_path(module_name: str) -> str:
    return str(Path(*module_name.split(".")).with_suffix(".py"))


def _state_key(extension_id: str) -> str:
    cleaned = re.sub(r"[^a-z0-9_]+", "_", extension_id.lower()).strip("_")
    return cleaned[:36] or "extension"


def _message_label(title: str) -> str:
    words = re.findall(r"[A-Za-z0-9]+", title.upper())
    return " ".join(words[:3]) if words else "EXTENSION"


def _primary_word(brief: Any) -> str:
    terms = _semantic_terms(brief)
    return terms[0] if terms else "phase"


def _semantic_terms(brief: Any) -> list[str]:
    text = " ".join((brief.extension_id, brief.title, brief.player_facing_goal, *brief.implementation_notes)).lower()
    candidates = (
        "blocker",
        "zapper",
        "slow",
        "shield",
        "overcharge",
        "safe_zone",
        "poison",
        "heavy",
        "phase",
        "charge",
        "cooldown",
        "streak",
        "alarm",
        "stealth",
        "extract",
        "tower",
        "relay",
        "trap",
        "mirror",
        "drone",
        "gate_type",
    )
    found = [term for term in candidates if term.replace("_", " ") in text or term in text]
    return found or ["phase", "charge", "cooldown"]


if __name__ == "__main__":
    main()