vibe/tests/core/test_rewind_manager.py

from __future__ import annotations

from pathlib import Path
from unittest.mock import patch

import pytest

from vibe.core.rewind.manager import FileSnapshot, RewindError, RewindManager
from vibe.core.types import LLMMessage, MessageList, Role


def _make_messages(*contents: str) -> MessageList:
    """Create a MessageList with a system message followed by user/assistant pairs."""
    msgs = MessageList([LLMMessage(role=Role.system, content="system")])
    for content in contents:
        msgs.append(LLMMessage(role=Role.user, content=content))
        msgs.append(LLMMessage(role=Role.assistant, content=f"reply to {content}"))
    return msgs


def _snap(path: Path) -> FileSnapshot:
    """Create a FileSnapshot by reading a file (or None if missing)."""
    resolved = str(path.resolve())
    try:
        content: bytes | None = path.read_bytes()
    except FileNotFoundError:
        content = None
    return FileSnapshot(path=resolved, content=content)


def _make_manager(
    messages: MessageList,
) -> tuple[RewindManager, list[bool], list[bool]]:
    save_calls: list[bool] = []
    reset_calls: list[bool] = []

    async def save_messages() -> None:
        save_calls.append(True)

    async def reset_session() -> None:
        reset_calls.append(True)

    mgr = RewindManager(
        messages=messages, save_messages=save_messages, reset_session=reset_session
    )
    return mgr, save_calls, reset_calls


class TestCheckpoints:
    def test_create_checkpoint_carries_forward_snapshots(self, tmp_path: Path) -> None:
        messages = _make_messages("hello", "world")
        mgr, _, _ = _make_manager(messages)
        f = tmp_path / "f.txt"
        f.write_text("v1", encoding="utf-8")

        mgr.create_checkpoint()
        mgr.add_snapshot(_snap(f))
        f.write_text("v2", encoding="utf-8")

        mgr.create_checkpoint()

        # Second checkpoint should have re-read the file
        assert len(mgr.checkpoints) == 2
        assert mgr.checkpoints[1].files[0].content == b"v2"

    def test_add_snapshot_to_all_checkpoints(self, tmp_path: Path) -> None:
        messages = _make_messages("hello", "world")
        mgr, _, _ = _make_manager(messages)

        mgr.create_checkpoint()
        mgr.create_checkpoint()

        f = tmp_path / "late.txt"
        f.write_text("content", encoding="utf-8")
        mgr.add_snapshot(_snap(f))

        resolved = str(f.resolve())
        assert any(s.path == resolved for s in mgr.checkpoints[0].files)
        assert any(s.path == resolved for s in mgr.checkpoints[1].files)

    def test_add_snapshot_no_duplicate(self, tmp_path: Path) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)
        mgr.create_checkpoint()

        f = tmp_path / "f.txt"
        f.write_text("content", encoding="utf-8")
        mgr.add_snapshot(_snap(f))
        mgr.add_snapshot(_snap(f))

        resolved = str(f.resolve())
        matches = [s for s in mgr.checkpoints[0].files if s.path == resolved]
        assert len(matches) == 1

    def test_has_changes_detects_new_file(self, tmp_path: Path) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)
        f = tmp_path / "new.txt"

        mgr.create_checkpoint()
        mgr.add_snapshot(FileSnapshot(path=str(f.resolve()), content=None))
        assert not mgr.has_file_changes_at(len(messages))

        f.write_text("created", encoding="utf-8")
        assert mgr.has_file_changes_at(len(messages))

    def test_has_changes_false_when_unchanged(self, tmp_path: Path) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)
        f = tmp_path / "f.txt"
        f.write_text("content", encoding="utf-8")

        mgr.create_checkpoint()
        mgr.add_snapshot(_snap(f))
        assert not mgr.has_file_changes_at(len(messages))

    def test_has_file_changes_at_no_checkpoint(self) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)
        assert not mgr.has_file_changes_at(1)


class TestRewind:
    def test_get_rewindable_messages(self) -> None:
        messages = _make_messages("hello", "world")
        mgr, _, _ = _make_manager(messages)

        result = mgr.get_rewindable_messages()

        assert len(result) == 2
        assert result[0] == (1, "hello")
        assert result[1] == (3, "world")

    def test_get_rewindable_messages_excludes_injected(self) -> None:
        messages = _make_messages("hello")
        # Insert an injected middleware message between turns
        messages.append(
            LLMMessage(role=Role.user, content="plan mode reminder", injected=True)
        )
        messages.append(LLMMessage(role=Role.user, content="world"))
        messages.append(LLMMessage(role=Role.assistant, content="reply to world"))
        mgr, _, _ = _make_manager(messages)

        result = mgr.get_rewindable_messages()

        assert len(result) == 2
        assert result[0] == (1, "hello")
        # Index 3 is the injected message — it must be skipped
        assert result[1] == (4, "world")

    @pytest.mark.asyncio
    async def test_rewind_to_message(self) -> None:
        messages = _make_messages("hello", "world")
        mgr, save_calls, reset_calls = _make_manager(messages)

        content, errors = await mgr.rewind_to_message(3, restore_files=False)

        assert content == "world"
        assert errors == []
        assert len(save_calls) == 1
        assert len(reset_calls) == 1
        assert len(messages) == 3

    @pytest.mark.asyncio
    async def test_rewind_to_message_invalid_index(self) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)

        with pytest.raises(RewindError, match="Invalid message index"):
            await mgr.rewind_to_message(99, restore_files=False)

    @pytest.mark.asyncio
    async def test_rewind_to_message_not_user(self) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)

        with pytest.raises(RewindError, match="not a user message"):
            await mgr.rewind_to_message(2, restore_files=False)

    def test_messages_reset_clears_checkpoints(self) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)
        mgr.create_checkpoint()

        assert len(mgr.checkpoints) == 1

        messages.reset([LLMMessage(role=Role.system, content="system")])

        assert len(mgr.checkpoints) == 0

    @pytest.mark.asyncio
    async def test_rewind_preserves_earlier_checkpoints(self) -> None:
        messages = _make_messages("hello", "world")
        mgr, _, _ = _make_manager(messages)

        mgr.create_checkpoint()
        mgr._checkpoints[-1].message_index = 1
        mgr.create_checkpoint()
        mgr._checkpoints[-1].message_index = 3

        assert len(mgr.checkpoints) == 2

        await mgr.rewind_to_message(3, restore_files=False)

        assert len(mgr.checkpoints) == 1

    def test_update_system_prompt_preserves_checkpoints(self) -> None:
        """Switching agents via shift+tab calls update_system_prompt which must
        NOT clear rewind checkpoints (unlike a full reset).
        """
        messages = _make_messages("hello", "world")
        mgr, _, _ = _make_manager(messages)

        mgr.create_checkpoint()
        mgr._checkpoints[-1].message_index = 1
        mgr.create_checkpoint()
        mgr._checkpoints[-1].message_index = 3

        assert len(mgr.checkpoints) == 2

        # Simulate shift+tab agent switch: only the system prompt changes
        messages.update_system_prompt("new agent system prompt")

        assert len(mgr.checkpoints) == 2
        assert messages[0].content == "new agent system prompt"

    def test_create_checkpoint_uses_current_message_count(self) -> None:
        messages = _make_messages("hello")
        mgr, _, _ = _make_manager(messages)

        mgr.create_checkpoint()

        assert mgr.checkpoints[0].message_index == len(messages)


class _Turn:
    """Helper that simulates one conversation turn."""

    def __init__(self, mgr: RewindManager, messages: MessageList) -> None:
        self._mgr = mgr
        self._messages = messages

    def begin(self, user_msg: str) -> None:
        """Start a new turn: create checkpoint then append user message.

        This mirrors agent_loop.act() which calls create_checkpoint()
        *before* the user message is added to the message list.
        """
        self._mgr.create_checkpoint()
        self._messages.append(LLMMessage(role=Role.user, content=user_msg))

    def tool_write(self, path: Path, content: str) -> None:
        """Simulate a tool writing to a file (snapshot → write)."""
        self._mgr.add_snapshot(_snap(path))
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8")

    def tool_delete(self, path: Path) -> None:
        """Simulate a tool deleting a file (snapshot → unlink)."""
        self._mgr.add_snapshot(_snap(path))
        path.unlink()

    def end(self, assistant_reply: str = "ok") -> None:
        """End the turn: append assistant reply."""
        self._messages.append(LLMMessage(role=Role.assistant, content=assistant_reply))


@pytest.mark.asyncio
class TestRewindScenarios:
    @staticmethod
    def _setup() -> tuple[RewindManager, MessageList, _Turn]:
        messages = MessageList([LLMMessage(role=Role.system, content="system")])
        mgr, _, _ = _make_manager(messages)
        return mgr, messages, _Turn(mgr, messages)

    async def test_edit_file_across_turns_rewind_to_middle(
        self, tmp_path: Path
    ) -> None:
        """A file is edited every turn. Rewinding restores the version from
        *before* the target turn.
        """
        mgr, messages, turn = self._setup()
        f = tmp_path / "app.py"
        f.write_text("v0", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, "v1")
        turn.end()

        turn.begin("turn2")
        turn.tool_write(f, "v2")
        turn.end()

        turn.begin("turn3")
        turn.tool_write(f, "v3")
        turn.end()

        # Rewind to turn2 (message index 3) → file should be v1
        rewindable = mgr.get_rewindable_messages()
        turn2_idx = rewindable[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)

        assert f.read_text(encoding="utf-8") == "v1"

    async def test_file_created_then_rewind_before_creation(
        self, tmp_path: Path
    ) -> None:
        """A file that didn't exist before should be deleted on rewind."""
        mgr, messages, turn = self._setup()
        new_file = tmp_path / "generated.py"

        turn.begin("turn1")
        turn.end()

        turn.begin("turn2")
        turn.tool_write(new_file, "print('hello')")
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]
        await mgr.rewind_to_message(turn1_idx, restore_files=True)

        assert not new_file.exists()

    async def test_file_deleted_by_tool_rewind_restores(self, tmp_path: Path) -> None:
        """Rewinding past a deletion restores the file."""
        mgr, _, turn = self._setup()
        f = tmp_path / "config.yaml"
        f.write_text("key: value", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, "key: value")  # touch to start tracking
        turn.end()

        turn.begin("turn2")
        turn.tool_delete(f)
        turn.end()

        assert not f.exists()

        turn2_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)

        assert f.exists()
        assert f.read_text(encoding="utf-8") == "key: value"

    async def test_mixed_create_and_edit(self, tmp_path: Path) -> None:
        """Multiple files: one pre-existing and edited, one created mid-session."""
        mgr, messages, turn = self._setup()
        existing = tmp_path / "main.py"
        existing.write_text("original", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(existing, "modified")
        turn.end()

        turn.begin("turn2")
        new_file = tmp_path / "utils.py"
        turn.tool_write(new_file, "def helper(): ...")
        turn.tool_write(existing, "modified again")
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]
        await mgr.rewind_to_message(turn1_idx, restore_files=True)

        assert existing.read_text(encoding="utf-8") == "original"
        assert not new_file.exists()

    async def test_user_manual_edit_between_turns(self, tmp_path: Path) -> None:
        """If the user edits a file between turns, the checkpoint at the next
        turn captures the user's version, so rewind restores it.
        """
        mgr, _, turn = self._setup()
        f = tmp_path / "readme.md"
        f.write_text("initial", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, "tool wrote this")
        turn.end()

        # User manually edits the file outside the tool loop
        f.write_text("user edited this", encoding="utf-8")

        turn.begin("turn2")
        turn.tool_write(f, "tool overwrote user")
        turn.end()

        # Rewind to turn2 → should restore the user's manual edit
        turn2_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)

        assert f.read_text(encoding="utf-8") == "user edited this"

    async def test_rewind_without_restore(self, tmp_path: Path) -> None:
        """Rewinding with restore_files=False truncates messages but keeps
        files as they are.
        """
        mgr, messages, turn = self._setup()
        f = tmp_path / "data.json"
        f.write_text("{}", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, '{"a": 1}')
        turn.end()

        turn.begin("turn2")
        turn.tool_write(f, '{"a": 1, "b": 2}')
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]
        await mgr.rewind_to_message(turn1_idx, restore_files=False)

        # File untouched
        assert f.read_text(encoding="utf-8") == '{"a": 1, "b": 2}'
        # But messages were truncated
        assert len(messages) == 1  # only system message

    async def test_rewind_then_new_turns_then_rewind_again(
        self, tmp_path: Path
    ) -> None:
        """After a rewind, new turns create new checkpoints. A second rewind
        should work correctly with the new history.
        """
        mgr, _, turn = self._setup()
        f = tmp_path / "code.py"
        f.write_text("v0", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, "v1")
        turn.end()

        turn.begin("turn2")
        turn.tool_write(f, "v2")
        turn.end()

        # Rewind to turn2
        turn2_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)
        assert f.read_text(encoding="utf-8") == "v1"

        # New turn after rewind
        turn.begin("turn2-bis")
        turn.tool_write(f, "v2-bis")
        turn.end()

        turn.begin("turn3-bis")
        turn.tool_write(f, "v3-bis")
        turn.end()

        # Rewind to turn2-bis
        turn2bis_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2bis_idx, restore_files=True)
        assert f.read_text(encoding="utf-8") == "v1"

    async def test_agent_switch_between_turns_preserves_rewind(
        self, tmp_path: Path
    ) -> None:
        """Pressing shift+tab between two messages switches agents, which calls
        update_system_prompt.  Checkpoints must survive so a subsequent rewind
        restores files correctly.
        """
        mgr, messages, turn = self._setup()
        f = tmp_path / "main.py"
        f.write_text("v0", encoding="utf-8")

        turn.begin("turn1")
        turn.tool_write(f, "v1")
        turn.end()

        # User presses shift+tab → agent switch → system prompt replaced
        messages.update_system_prompt("switched agent prompt")

        turn.begin("turn2")
        turn.tool_write(f, "v2")
        turn.end()

        # Rewind to turn2 should restore "v1"
        turn2_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)

        assert f.read_text(encoding="utf-8") == "v1"

    async def test_binary_file_snapshot_and_restore(self, tmp_path: Path) -> None:
        """Binary files (non-UTF-8) are snapshotted and restored correctly."""
        mgr, _, turn = self._setup()
        f = tmp_path / "image.bin"
        original = bytes(range(256))
        f.write_bytes(original)

        turn.begin("turn1")
        mgr.add_snapshot(_snap(f))
        f.write_bytes(b"\x00" * 256)
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]
        await mgr.rewind_to_message(turn1_idx, restore_files=True)

        assert f.read_bytes() == original

    async def test_create_edit_delete_full_lifecycle(self, tmp_path: Path) -> None:
        """File goes through create → edit → delete. Rewind to each point
        restores the correct state.
        """
        mgr, _, turn = self._setup()
        f = tmp_path / "temp.txt"

        turn.begin("turn1")
        turn.tool_write(f, "created")
        turn.end()

        turn.begin("turn2")
        turn.tool_write(f, "edited")
        turn.end()

        turn.begin("turn3")
        turn.tool_delete(f)
        turn.end()

        assert not f.exists()

        # Rewind to turn3 → file should be "edited" (state before deletion)
        turn3_idx = mgr.get_rewindable_messages()[2][0]
        await mgr.rewind_to_message(turn3_idx, restore_files=True)
        assert f.read_text(encoding="utf-8") == "edited"

    async def test_user_creates_file_tool_overwrites(self, tmp_path: Path) -> None:
        """User creates a file manually before a turn. The tool overwrites it.
        Rewind restores the user's version.
        """
        mgr, _, turn = self._setup()
        f = tmp_path / "notes.txt"

        turn.begin("turn1")
        turn.end()

        # User creates the file manually between turns
        f.write_text("user notes", encoding="utf-8")

        turn.begin("turn2")
        turn.tool_write(f, "overwritten by tool")
        turn.end()

        turn2_idx = mgr.get_rewindable_messages()[1][0]
        await mgr.rewind_to_message(turn2_idx, restore_files=True)

        assert f.read_text(encoding="utf-8") == "user notes"

    async def test_nested_directory_files(self, tmp_path: Path) -> None:
        """Files in nested directories are restored including parent dirs."""
        mgr, _, turn = self._setup()
        deep = tmp_path / "src" / "pkg" / "module.py"

        turn.begin("turn1")
        turn.tool_write(deep, "def foo(): pass")
        turn.end()

        turn.begin("turn2")
        turn.tool_write(deep, "def foo(): return 42")
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]

        # Delete everything
        deep.unlink()
        (tmp_path / "src" / "pkg").rmdir()
        (tmp_path / "src").rmdir()

        await mgr.rewind_to_message(turn1_idx, restore_files=True)

        # File didn't exist before turn1 → should be deleted
        assert not deep.exists()

    async def test_rewind_restores_errors_collected(self, tmp_path: Path) -> None:
        """When removing a file during rewind fails, errors are returned in the tuple."""
        mgr, _, turn = self._setup()
        created_file = tmp_path / "locked.txt"

        turn.begin("turn1")
        turn.end()

        turn.begin("turn2")
        # Snapshot runs before write → earlier checkpoints record content=None
        turn.tool_write(created_file, "created in turn2")
        turn.end()

        turn1_idx = mgr.get_rewindable_messages()[0][0]
        with patch(
            "vibe.core.rewind.manager.os.remove",
            side_effect=OSError("mocked removal failure"),
        ):
            _, errors = await mgr.rewind_to_message(turn1_idx, restore_files=True)

        assert len(errors) == 1
        assert "Failed to delete file" in errors[0]
        assert "locked.txt" in errors[0]