tests/test_code_mutation.py

import json
import importlib.util
import sys
from pathlib import Path
from types import SimpleNamespace

import pytest

from life_game.code_mutation import (
    CodeMutationSpecError,
    apply_method_replacement_spec,
    apply_code_mutation_spec,
    apply_unified_diff_patch,
    code_mutation_diff_stats,
    code_mutation_spec_to_dict,
    extract_unified_diff_patch,
    method_replacement_spec_to_dict,
    parse_code_mutation_spec,
    parse_method_replacement_spec,
    parse_unified_diff_patch_spec,
    render_method_replacement_diff,
    render_code_mutation_diff,
)
from life_game.code_mutation_runtime import apply_method_replacements_live, select_runtime_extension
from life_game.code_extensions import (
    EXTENSION_BRIEFS_BY_MODE,
    extension_briefs_for_mode,
    format_extension_briefs,
    select_extension_brief,
)
from life_game.game import FIREWALL_MODE, GAME_MODES, SANDBOX_MODE, GameModel
from life_game.games.firewall import FirewallGame


def _load_propose_code_mutation_module():
    path = Path("scripts/propose_code_mutation.py")
    spec = importlib.util.spec_from_file_location("propose_code_mutation", path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _load_code_mutation_sft_seed_module():
    path = Path("scripts/generate_code_mutation_sft_seed.py")
    spec = importlib.util.spec_from_file_location("generate_code_mutation_sft_seed", path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _spec(path="life_game/games/data_vampire.py", old_lines=None, new_lines=None):
    if old_lines is None:
        old_lines = ["old"]
    if new_lines is None:
        new_lines = ["new"]
    return json.dumps(
        {
            "summary": "Tune data vampire",
            "rationale": "The mode needs clearer food behavior.",
            "edits": [
                {
                    "path": path,
                    "purpose": "Update the selected game behavior.",
                    "old_lines": old_lines,
                    "new_lines": new_lines,
                }
            ],
            "tests": ["uv run pytest tests/test_game.py"],
        }
    )


def _range_spec(path="life_game/games/data_vampire.py", start_line=2, end_line=3, new_lines=None):
    if new_lines is None:
        new_lines = ["replacement"]
    return json.dumps(
        {
            "summary": "Tune data vampire",
            "rationale": "The mode needs clearer food behavior.",
            "edits": [
                {
                    "path": path,
                    "purpose": "Update a selected line range.",
                    "start_line": start_line,
                    "end_line": end_line,
                    "new_lines": new_lines,
                }
            ],
            "tests": ["uv run pytest tests/test_game.py"],
        }
    )


def test_parse_code_mutation_spec_accepts_allowed_game_edit():
    spec = parse_code_mutation_spec(_spec())

    assert spec.summary == "Tune data vampire"
    assert spec.edits[0].path == "life_game/games/data_vampire.py"
    assert spec.tests == ("uv run pytest tests/test_game.py",)
    assert code_mutation_spec_to_dict(spec)["edits"][0]["purpose"] == "Update the selected game behavior."
    assert spec.edits[0].old == "old\n"
    assert spec.edits[0].new == "new\n"


def test_parse_code_mutation_spec_accepts_line_range_edit():
    spec = parse_code_mutation_spec(_range_spec(start_line=2, end_line=4))

    assert spec.edits[0].start_line == 2
    assert spec.edits[0].end_line == 4


def test_parse_method_replacement_spec_ignores_optional_line_metadata():
    spec = parse_method_replacement_spec(
        json.dumps(
            {
                "summary": "Replace method",
                "rationale": "The method changes the game behavior.",
                "replacements": [
                    {
                        "path": "life_game/games/data_vampire.py",
                        "class_name": "DataVampireGame",
                        "method_name": "advance",
                        "content": "    def advance(self, model, game, rule):\n        return model, game\n",
                        "start_line": 10,
                        "end_line": 12,
                    }
                ],
                "tests": ["uv run pytest tests/test_game.py"],
            }
        )
    )

    assert spec.replacements[0].method_name == "advance"


def test_parse_method_replacement_spec_accepts_long_qwen_summary():
    spec = parse_method_replacement_spec(
        json.dumps(
            {
                "summary": "Implement tactical anchor modes for drone formations by storing state in game data and modifying command and advance to support guard and orbit behaviors.",
                "rationale": "The method changes the game behavior.",
                "replacements": [
                    {
                        "path": "life_game/games/data_vampire.py",
                        "class_name": "DataVampireGame",
                        "method_name": "advance",
                        "content": "    def advance(self, model, game, rule):\n        return model, game\n",
                    }
                ],
                "tests": ["uv run pytest tests/test_game.py"],
            }
        )
    )

    assert spec.summary.startswith("Implement tactical anchor modes")


def test_parse_method_replacement_spec_splits_duplicate_method_pairs_with_shared_path():
    raw = """{
      "summary": "Replace drone methods",
      "rationale": "The replacements alter drone behavior.",
      "replacements": [
        {
          "path": "life_game/games/drone_swarm.py",
          "class_name": "DroneSwarmGame",
          "method_name": "command",
          "content": "def command(self, model, game, command):\\n    return model, game\\n",
          "method_name": "advance",
          "content": "def advance(self, model, game, rule):\\n    return model, game\\n"
        }
      ],
      "tests": ["uv run pytest tests/test_game.py"]
    }"""

    spec = parse_method_replacement_spec(raw)

    assert [replacement.method_name for replacement in spec.replacements] == ["command", "advance"]
    assert all(replacement.path == "life_game/games/drone_swarm.py" for replacement in spec.replacements)


def test_select_runtime_extension_is_deterministic():
    first = select_runtime_extension("Firewall Defender", seed=0)
    repeated = select_runtime_extension("Firewall Defender", seed=0)
    alternate = select_runtime_extension("Firewall Defender", seed=1)

    assert first.extension_id == repeated.extension_id
    assert alternate.extension_id


def test_apply_method_replacements_live_patches_registered_game():
    original = FirewallGame.progress
    spec = parse_method_replacement_spec(
        json.dumps(
            {
                "summary": "Patch progress",
                "rationale": "The live registry should use validated replacement methods.",
                "replacements": [
                    {
                        "path": "life_game/games/firewall.py",
                        "class_name": "FirewallGame",
                        "method_name": "progress",
                        "content": (
                            "    def progress(self, game):\n"
                            "        return \"Patched\", \"live\", None\n"
                        ),
                    }
                ],
                "tests": ["uv run pytest tests/test_code_mutation.py"],
            }
        )
    )
    try:
        applied = apply_method_replacements_live(spec)

        assert applied == ("FirewallGame.progress",)
        assert FirewallGame().progress(GameModel(mode=FIREWALL_MODE)) == ("Patched", "live", None)
    finally:
        FirewallGame.progress = original


def test_apply_method_replacements_live_rejects_signature_mismatch():
    spec = parse_method_replacement_spec(
        json.dumps(
            {
                "summary": "Bad progress",
                "rationale": "The replacement drops required method arguments.",
                "replacements": [
                    {
                        "path": "life_game/games/firewall.py",
                        "class_name": "FirewallGame",
                        "method_name": "progress",
                        "content": "    def progress(self):\n        return None\n",
                    }
                ],
                "tests": ["uv run pytest tests/test_code_mutation.py"],
            }
        )
    )

    with pytest.raises(CodeMutationSpecError, match="signature"):
        apply_method_replacements_live(spec)


def test_repair_prompt_preserves_method_replacement_contract():
    module = _load_propose_code_mutation_module()

    messages = module._repair_prompt_messages("method-replacements", "{}", "bad schema")

    assert '"replacements"' in messages[1]["content"]
    assert '"edits"' not in messages[1]["content"]
    assert "full corrected proposal" in messages[1]["content"]


def test_method_replacement_prompt_uses_compact_relevant_source():
    module = _load_propose_code_mutation_module()

    messages = module._prompt_messages(
        "Drone Swarm",
        "Make formation behavior tactical and deterministic.",
        common_chars=1200,
        extension_id="anchor_modes",
        ambitious=True,
        min_changed_lines=30,
        proposal_format="method-replacements",
    )
    text = "\n".join(message["content"] for message in messages)

    assert len(text) < 10000
    assert "Relevant numbered excerpts from life_game/games/drone_swarm.py" in text
    assert "Required method signatures:" in text
    assert "def command(self, model: BoardModel, game: GameModel, command: str)" in text
    assert "def advance(self, model: BoardModel, game: GameModel, rule: LifeRule)" in text
    assert "def _drones" not in text
    assert "Replace only these required methods" in text


def test_code_mutation_sft_seed_rows_are_sharegpt_and_parseable():
    module = _load_code_mutation_sft_seed_module()

    rows = module.build_dataset(12)

    assert len(rows) == 12
    assert {row["schema"] for row in rows} == {"signal-garden-code-mutation-sft/v1"}
    for row in rows:
        assert [message["from"] for message in row["conversations"]] == ["system", "human", "gpt"]
        assert "Do not invent helper methods" in row["conversations"][1]["value"]
        assert parse_method_replacement_spec(row["conversations"][-1]["value"]).replacements


def test_gatekeeper_typed_gate_semantic_validation_rejects_noop():
    module = _load_propose_code_mutation_module()
    method_spec = SimpleNamespace(
        replacements=(
            SimpleNamespace(method_name="new", content="def new(self, size, rng):\n    data = {'gates': ()}\n"),
            SimpleNamespace(method_name="command", content="def command(self, model, game, command):\n    return model, game\n"),
            SimpleNamespace(method_name="advance", content="def advance(self, model, game, rule):\n    return model, game\n"),
        )
    )

    with pytest.raises(RuntimeError, match="gate_type"):
        module._validate_extension_method_content(method_spec, "Gatekeeper", "typed_gates")


def test_gatekeeper_typed_gate_semantic_requirement_is_prompted():
    module = _load_propose_code_mutation_module()

    text = module._semantic_requirement_text("Gatekeeper", "typed_gates")

    assert "gate_type" in text
    assert "blocker" in text
    assert "zapper" in text
    assert "slow" in text


def test_parse_code_mutation_spec_wraps_top_level_edit_object():
    spec = parse_code_mutation_spec(
        json.dumps(
            {
                "path": "life_game/games/data_vampire.py",
                "purpose": "Update a selected line range.",
                "start_line": 2,
                "end_line": 3,
                "new_lines": ["replacement"],
            }
        )
    )

    assert spec.summary == "Apply code edit"
    assert spec.edits[0].start_line == 2


def test_parse_code_mutation_spec_wraps_top_level_edit_list():
    spec = parse_code_mutation_spec(
        json.dumps(
            [
                {
                    "path": "life_game/games/data_vampire.py",
                    "purpose": "Update a selected line range.",
                    "start_line": 2,
                    "end_line": 3,
                    "new_lines": ["replacement"],
                }
            ]
        )
    )

    assert spec.summary == "Apply code edits"
    assert len(spec.edits) == 1


def test_parse_code_mutation_spec_extracts_prefixed_json():
    spec = parse_code_mutation_spec(f"</think>{_spec()}\nextra text")

    assert spec.edits[0].path == "life_game/games/data_vampire.py"


def test_parse_code_mutation_spec_rejects_unsafe_paths():
    with pytest.raises(CodeMutationSpecError, match="not allowed"):
        parse_code_mutation_spec(_spec(path="app.py"))

    with pytest.raises(CodeMutationSpecError, match="Unsafe"):
        parse_code_mutation_spec(_spec(path="../life_game/games/data_vampire.py"))


def test_apply_code_mutation_spec_writes_only_when_not_dry_run(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("old\n", encoding="utf-8")
    spec = parse_code_mutation_spec(_spec(old_lines=["old"], new_lines=["new"]))

    planned = apply_code_mutation_spec(spec, tmp_path, dry_run=True)
    assert planned == (target.resolve(),)
    assert target.read_text(encoding="utf-8") == "old\n"

    written = apply_code_mutation_spec(spec, tmp_path, dry_run=False)
    assert written == (target.resolve(),)
    assert target.read_text(encoding="utf-8") == "new\n"


def test_apply_code_mutation_spec_rejects_missing_or_ambiguous_old_block(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("same\nsame\n", encoding="utf-8")

    ambiguous = parse_code_mutation_spec(_spec(old_lines=["same"], new_lines=["different"]))
    with pytest.raises(CodeMutationSpecError, match="exactly once; found 2"):
        apply_code_mutation_spec(ambiguous, tmp_path)

    missing = parse_code_mutation_spec(_spec(old_lines=["absent"], new_lines=["different"]))
    with pytest.raises(CodeMutationSpecError, match="exactly once; found 0"):
        apply_code_mutation_spec(missing, tmp_path)


def test_render_code_mutation_diff_shows_exact_edit(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("old\nkeep\n", encoding="utf-8")
    spec = parse_code_mutation_spec(_spec(old_lines=["old"], new_lines=["new"]))

    diff = render_code_mutation_diff(spec, tmp_path)

    assert "--- a/life_game/games/data_vampire.py" in diff
    assert "+++ b/life_game/games/data_vampire.py" in diff
    assert "-old\n" in diff
    assert "+new\n" in diff


def test_code_mutation_diff_stats_counts_changed_lines():
    diff = "\n".join(
        (
            "--- a/life_game/games/firewall.py",
            "+++ b/life_game/games/firewall.py",
            "@@ -1,2 +1,3 @@",
            " keep",
            "-old",
            "+new",
            "+extra",
        )
    )

    stats = code_mutation_diff_stats(diff)

    assert stats.added_lines == 2
    assert stats.removed_lines == 1
    assert stats.changed_lines == 3


def test_extract_unified_diff_patch_from_fenced_response():
    raw = """Here is the patch:
```diff
diff --git a/life_game/games/data_vampire.py b/life_game/games/data_vampire.py
--- a/life_game/games/data_vampire.py
+++ b/life_game/games/data_vampire.py
@@ -1 +1 @@
-old
+new
```
"""

    patch = extract_unified_diff_patch(raw)

    assert patch.startswith("diff --git")
    assert "+new" in patch


def test_parse_unified_diff_patch_spec_accepts_json_patch_wrapper():
    patch = """diff --git a/life_game/games/data_vampire.py b/life_game/games/data_vampire.py
--- a/life_game/games/data_vampire.py
+++ b/life_game/games/data_vampire.py
@@ -1 +1 @@
-old
+new
"""
    spec = parse_unified_diff_patch_spec(
        json.dumps(
            {
                "summary": "Patch data vampire",
                "rationale": "The mode needs clearer behavior.",
                "patch": patch,
                "tests": ["uv run pytest tests/test_game.py"],
            }
        )
    )

    assert spec["summary"] == "Patch data vampire"
    assert spec["patch"].startswith("diff --git")
    assert spec["tests"] == ["uv run pytest tests/test_game.py"]


def test_parse_unified_diff_patch_spec_defaults_optional_metadata():
    patch = """diff --git a/life_game/games/data_vampire.py b/life_game/games/data_vampire.py
--- a/life_game/games/data_vampire.py
+++ b/life_game/games/data_vampire.py
@@ -1 +1 @@
-old
+new
"""
    spec = parse_unified_diff_patch_spec(json.dumps({"patch": patch}))

    assert spec["summary"] == "Apply unified diff patch"
    assert spec["tests"] == ["uv run pytest tests/test_game.py"]


def test_method_replacement_spec_replaces_complete_method(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text(
        "class DataVampireGame:\n"
        "    def command(self):\n"
        "        return 'old'\n"
        "\n"
        "    def advance(self):\n"
        "        return 'keep'\n",
        encoding="utf-8",
    )
    raw = json.dumps(
        {
            "summary": "Patch data vampire",
            "rationale": "Replace one method.",
            "replacements": [
                {
                    "path": "life_game/games/data_vampire.py",
                    "class_name": "DataVampireGame",
                    "method_name": "command",
                    "content": "def command(self):\n    return 'new'",
                }
            ],
            "tests": ["uv run pytest tests/test_game.py"],
        }
    )
    spec = parse_method_replacement_spec(raw)

    assert method_replacement_spec_to_dict(spec)["replacements"][0]["method_name"] == "command"
    diff = render_method_replacement_diff(spec, tmp_path)
    assert "-        return 'old'\n" in diff
    assert "+        return 'new'\n" in diff

    planned = apply_method_replacement_spec(spec, tmp_path, dry_run=True)
    assert planned == (target.resolve(),)
    assert "return 'old'" in target.read_text(encoding="utf-8")

    apply_method_replacement_spec(spec, tmp_path, dry_run=False)
    assert "return 'new'" in target.read_text(encoding="utf-8")


def test_method_replacement_spec_splits_duplicate_replacement_keys():
    raw = """
{
  "replacements": [
    {
      "path": "life_game/games/data_vampire.py",
      "class_name": "DataVampireGame",
      "method_name": "command",
      "content": "def command(self):\\n    return 'new'",
      "path": "life_game/games/data_vampire.py",
      "class_name": "DataVampireGame",
      "method_name": "advance",
      "content": "def advance(self):\\n    return 'new'"
    }
  ]
}
"""

    spec = parse_method_replacement_spec(raw)

    assert [replacement.method_name for replacement in spec.replacements] == ["command", "advance"]


def test_method_replacement_spec_rejects_invalid_method_source(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("class DataVampireGame:\n    def command(self):\n        return 'old'\n", encoding="utf-8")
    spec = parse_method_replacement_spec(
        json.dumps(
            {
                "replacements": [
                    {
                        "path": "life_game/games/data_vampire.py",
                        "class_name": "DataVampireGame",
                        "method_name": "command",
                        "content": "def command(self):\n    if True",
                    }
                ]
            }
        )
    )

    with pytest.raises(CodeMutationSpecError, match="does not parse"):
        apply_method_replacement_spec(spec, tmp_path)


def test_apply_unified_diff_patch_writes_only_when_not_dry_run(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("old\nkeep\n", encoding="utf-8")
    patch = """diff --git a/life_game/games/data_vampire.py b/life_game/games/data_vampire.py
--- a/life_game/games/data_vampire.py
+++ b/life_game/games/data_vampire.py
@@ -1,2 +1,2 @@
-old
+new
 keep
"""

    planned = apply_unified_diff_patch(patch, tmp_path, dry_run=True)
    assert planned == (target.resolve(),)
    assert target.read_text(encoding="utf-8") == "old\nkeep\n"

    written = apply_unified_diff_patch(patch, tmp_path, dry_run=False)
    assert written == (target.resolve(),)
    assert target.read_text(encoding="utf-8") == "new\nkeep\n"


def test_apply_unified_diff_patch_rejects_unsafe_or_deleting_paths(tmp_path):
    unsafe = """diff --git a/app.py b/app.py
--- a/app.py
+++ b/app.py
@@ -1 +1 @@
-old
+new
"""
    with pytest.raises(CodeMutationSpecError, match="not allowed"):
        apply_unified_diff_patch(unsafe, tmp_path)

    delete = """diff --git a/life_game/games/data_vampire.py b/life_game/games/data_vampire.py
--- a/life_game/games/data_vampire.py
+++ /dev/null
@@ -1 +0,0 @@
-old
"""
    with pytest.raises(CodeMutationSpecError, match="delete"):
        apply_unified_diff_patch(delete, tmp_path)


def test_parse_code_mutation_spec_splits_lines_with_embedded_newlines():
    spec = parse_code_mutation_spec(_spec(old_lines=["old\nblock"], new_lines=["new\nblock"]))

    assert spec.edits[0].old_lines == ("old", "block")
    assert spec.edits[0].new_lines == ("new", "block")


def test_apply_code_mutation_spec_replaces_line_range(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("one\ntwo\nthree\nfour\n", encoding="utf-8")
    spec = parse_code_mutation_spec(_range_spec(start_line=2, end_line=3, new_lines=["TWO", "THREE"]))

    apply_code_mutation_spec(spec, tmp_path, dry_run=False)

    assert target.read_text(encoding="utf-8") == "one\nTWO\nTHREE\nfour\n"


def test_apply_code_mutation_spec_applies_multiple_line_ranges_against_original_file(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("one\ntwo\nthree\nfour\nfive\nsix\n", encoding="utf-8")
    spec = parse_code_mutation_spec(
        json.dumps(
            {
                "summary": "Tune data vampire",
                "rationale": "The mode needs clearer food behavior.",
                "edits": [
                    {
                        "path": "life_game/games/data_vampire.py",
                        "purpose": "Remove an early line.",
                        "start_line": 2,
                        "end_line": 2,
                        "new_lines": [],
                    },
                    {
                        "path": "life_game/games/data_vampire.py",
                        "purpose": "Replace later lines using original numbering.",
                        "start_line": 5,
                        "end_line": 6,
                        "new_lines": ["FIVE"],
                    },
                ],
                "tests": ["uv run pytest tests/test_game.py"],
            }
        )
    )

    apply_code_mutation_spec(spec, tmp_path, dry_run=False)

    assert target.read_text(encoding="utf-8") == "one\nthree\nfour\nFIVE\n"


def test_apply_code_mutation_spec_rejects_overlapping_line_ranges(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("one\ntwo\nthree\n", encoding="utf-8")
    spec = parse_code_mutation_spec(
        json.dumps(
            {
                "summary": "Tune data vampire",
                "rationale": "The mode needs clearer food behavior.",
                "edits": [
                    {
                        "path": "life_game/games/data_vampire.py",
                        "purpose": "Replace first block.",
                        "start_line": 1,
                        "end_line": 2,
                        "new_lines": ["ONE"],
                    },
                    {
                        "path": "life_game/games/data_vampire.py",
                        "purpose": "Replace overlapping block.",
                        "start_line": 2,
                        "end_line": 3,
                        "new_lines": ["TWO"],
                    },
                ],
                "tests": ["uv run pytest tests/test_game.py"],
            }
        )
    )

    with pytest.raises(CodeMutationSpecError, match="overlap"):
        apply_code_mutation_spec(spec, tmp_path)


def test_apply_code_mutation_spec_rejects_line_range_outside_file(tmp_path):
    target = tmp_path / "life_game/games/data_vampire.py"
    target.parent.mkdir(parents=True)
    target.write_text("one\n", encoding="utf-8")
    spec = parse_code_mutation_spec(_range_spec(start_line=1, end_line=2))

    with pytest.raises(CodeMutationSpecError, match="exceeds file length"):
        apply_code_mutation_spec(spec, tmp_path)


def test_extension_briefs_include_material_gameplay_guidance():
    brief = select_extension_brief("Firewall Defender", "spread_charge")

    assert brief is not None
    assert brief.target_change_lines >= 30
    assert "advance" in brief.required_methods
    assert "charge" in brief.player_facing_goal.lower()
    assert "spread_charge" in format_extension_briefs("Firewall Defender", "spread_charge")


def test_annotation_extension_briefs_cover_registered_modes():
    missing = [mode for mode in GAME_MODES if mode != SANDBOX_MODE and mode not in EXTENSION_BRIEFS_BY_MODE]

    assert missing == []
    assert select_extension_brief("Gravity Well", "polarity_charge_heavies") is not None
    assert "heavy enemies" in format_extension_briefs("Gravity Well", "polarity_charge_heavies").lower()
    assert extension_briefs_for_mode("Gatekeeper")[0].target_change_lines >= 30