if_engine.py

"""
Thin embeddable wrapper around xyppy's Z-machine engine.

Captures game text in memory (no terminal) and accepts player commands
via a queue. This is our glue layer — not part of xyppy itself.
"""

from __future__ import annotations

import argparse
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

from xyppy import ops, ops_impl, quetzal, zenv
from xyppy import vterm

from saves import resolve_save_path, save_file_path, set_active_story

# Patch quetzal once so in-game save/restore use games/saves/.
_orig_quetzal_read = quetzal.read
_orig_quetzal_write = quetzal.write


def _save_basename(filename: str) -> str:
    name = Path(filename.strip()).name
    if name.endswith(".sav"):
        name = name[:-4]
    return name or "save"


def _save_path(filename: str) -> Path:
    return save_file_path(filename)


def _read_save_path(filename: str) -> Path:
    path = resolve_save_path(filename)
    return path if path is not None else save_file_path(filename)


def _patched_quetzal_read(filename: str):
    return _orig_quetzal_read(str(_read_save_path(filename)))


def _patched_quetzal_write(env, filename: str) -> bool:
    return _orig_quetzal_write(env, str(_save_path(filename)))


quetzal.read = _patched_quetzal_read
quetzal.write = _patched_quetzal_write


class NeedInputError(Exception):
    """Raised when the game waits for player input and none is queued."""

    def __init__(self, kind: str = "line", prefilled: str = ""):
        self.kind = kind
        self.prefilled = prefilled
        super().__init__(kind)


class EmbeddableEnv(zenv.Env):
    """zenv.Env that keeps our CapturingScreen across env.reset() (restore)."""

    def __init__(self, mem: bytes, options: argparse.Namespace) -> None:
        self._screen_factory: Callable[[zenv.Env], "CapturingScreen"] | None = None
        super().__init__(mem, options)

    def attach_capturing_screen(
        self, factory: Callable[[zenv.Env], "CapturingScreen"]
    ) -> None:
        self._screen_factory = factory
        self._install_capturing_screen()

    def _install_capturing_screen(self) -> None:
        screen = self._screen_factory(self)  # type: ignore[misc]
        self.screen = screen
        self.output_buffer[1] = screen

    def reset(self) -> None:
        bits_to_save = self.hdr.flags2 & 3
        factory = self._screen_factory
        super().reset()
        self.hdr.flags2 &= ~3
        self.hdr.flags2 |= bits_to_save
        self._screen_factory = factory
        if factory is not None:
            self._install_capturing_screen()


@dataclass
class _UndoSnapshot:
    mem: list[int]
    pc: int
    callstack: list[ops_impl.Frame]
    last_pc_branch_var: int | None
    last_pc_store_var: int | None


@dataclass
class GameSession:
    """Runs a Z-code story file with in-memory I/O."""

    story_path: Path
    _input_queue: list[str] = field(default_factory=list, repr=False)
    _char_remainder: str = field(default="", repr=False)
    _finished: bool = field(default=False, repr=False)
    _env: Optional[EmbeddableEnv] = field(default=None, repr=False)
    _rewind_pc: Optional[int] = field(default=None, repr=False)
    _undo: _UndoSnapshot | None = field(default=None, repr=False)
    _undo_restore: _UndoSnapshot | None = field(default=None, repr=False)
    _pending_input_kind: str | None = field(default=None, repr=False)

    def __post_init__(self) -> None:
        self.story_path = Path(self.story_path)
        if not self.story_path.is_file():
            raise FileNotFoundError(self.story_path)

    def _screen_factory(self, env: zenv.Env) -> "CapturingScreen":
        env.hdr.screen_width_units = 80
        env.hdr.screen_width_chars = 80
        return CapturingScreen(env, self._next_line_input, self._next_char_input)

    def _build_env(self) -> EmbeddableEnv:
        mem = self.story_path.read_bytes()
        options = argparse.Namespace(no_slow_scroll=True)
        env = EmbeddableEnv(mem, options)
        env.attach_capturing_screen(self._screen_factory)
        env._turnkeeper_session = self  # noqa: SLF001
        ops.setup_opcodes(env)
        self._register_undo_opcodes(env)
        return env

    def _register_undo_opcodes(self, env: EmbeddableEnv) -> None:
        session = self

        def save_undo(env, opinfo) -> None:
            if session._capture_undo():
                ops_impl.set_var(env, opinfo.store_var, 1)
            else:
                ops_impl.set_var(env, opinfo.store_var, 0)

        def restore_undo(env, opinfo) -> None:
            if session._restore_undo():
                ops_impl.set_var(env, opinfo.store_var, 2)
            else:
                ops_impl.set_var(env, opinfo.store_var, 0)

        ops.ext(9, save_undo, svar=True)
        ops.ext(10, restore_undo, svar=True)

    def _copy_frame(self, frame: ops_impl.Frame) -> ops_impl.Frame:
        return ops_impl.Frame(
            frame.return_addr,
            frame.num_args,
            list(frame.locals),
            frame.return_val_loc,
            list(frame.stack),
        )

    def _capture_undo(self) -> bool:
        env = self.env
        try:
            snap = _UndoSnapshot(
                mem=list(env.mem[: env.hdr.static_mem_base]),
                pc=env.pc,
                callstack=[self._copy_frame(f) for f in env.callstack],
                last_pc_branch_var=env.last_pc_branch_var,
                last_pc_store_var=env.last_pc_store_var,
            )
            # Preserve the prior snapshot — restore_undo reverts to the state
            # saved at the start of the last command (Balances skips save_undo
            # on the undo turn itself).
            self._undo_restore = self._undo
            self._undo = snap
            return True
        except Exception:  # noqa: BLE001
            self._undo_restore = None
            self._undo = None
            return False

    def _restore_undo(self) -> bool:
        target = self._undo_restore if self._undo_restore is not None else self._undo
        if target is None:
            return False
        snap = target
        env = self.env
        for i, byte in enumerate(snap.mem):
            env.mem[i] = byte
        env.pc = snap.pc
        env.callstack = [self._copy_frame(f) for f in snap.callstack]
        env.last_pc_branch_var = snap.last_pc_branch_var
        env.last_pc_store_var = snap.last_pc_store_var
        env.fixup_after_restore()
        self.screen.reset_display()
        return True

    @property
    def env(self) -> EmbeddableEnv:
        if self._env is None:
            self._env = self._build_env()
        return self._env

    @property
    def screen(self) -> "CapturingScreen":
        return self.env.screen  # type: ignore[return-value]

    @property
    def needs_char_input(self) -> bool:
        """True when the game waits for a single keystroke (press any key)."""
        return self._pending_input_kind == "char"

    def _next_line_input(self, prefilled: str = "") -> str:
        if self._char_remainder:
            word = self._char_remainder
            self._char_remainder = ""
            return word
        if self._input_queue:
            return self._input_queue.pop(0)
        if prefilled:
            return prefilled
        raise NeedInputError("line", prefilled)

    def _next_char_input(self) -> str:
        if self._char_remainder:
            ch = self._char_remainder[0]
            self._char_remainder = self._char_remainder[1:]
            return ch
        if self._input_queue:
            line = self._input_queue.pop(0)
            if line:
                self._char_remainder = line[1:]
                return line[0]
        raise NeedInputError("char")

    def _safe_step(self) -> None:
        """Run one VM step; rewind PC if input is needed mid-read."""
        if self._rewind_pc is not None:
            self.env.pc = self._rewind_pc
            self._rewind_pc = None
        pc_before = self.env.pc
        try:
            zenv.step(self.env)
        except NeedInputError:
            # xyppy advances PC before the read opcode runs; rewind so the
            # read retries once input is queued (avoids a phantom empty turn).
            self._rewind_pc = pc_before
            raise

    def reset(self) -> None:
        """Start a fresh game from the story file."""
        self._input_queue.clear()
        self._char_remainder = ""
        self._finished = False
        self._rewind_pc = None
        self._undo = None
        self._undo_restore = None
        self._pending_input_kind = None
        self._env = self._build_env()
        self.screen.first_draw()

    def start(self) -> str:
        """Run from a clean state until the game needs input or quits."""
        self.reset()
        return self.run_until_input()

    def send_command(self, command: str) -> str:
        """Send a player command and return new game output."""
        if self._finished:
            raise RuntimeError("Game has ended — call reset() to play again.")
        cmd = command.strip()
        if cmd.lower() == "undo":
            return self._undo_last_turn()
        if self._pending_input_kind == "char":
            # Press-any-key prompts: consume one character only so leftovers
            # (e.g. typing "ok") are not executed as the next command.
            self._char_remainder = ""
            self._input_queue.append(cmd[0] if cmd else " ")
        else:
            self._input_queue.append(cmd)
        return self.run_until_input()

    def send_key(self, char: str = " ") -> str:
        """Send a single keystroke for press-any-key / more-text prompts."""
        if self._finished:
            raise RuntimeError("Game has ended — call reset() to play again.")
        self._char_remainder = ""
        self._input_queue.append((char or " ")[:1])
        return self.run_until_input()

    def _undo_last_turn(self) -> str:
        """Restore the pre-turn snapshot and refresh the room description."""
        if self._undo_restore is None and self._undo is None:
            return '[You can\'t "undo" what hasn\'t been done!]'
        if not self._restore_undo():
            return '"Undo" failed. [Not all interpreters provide it.]'
        self.screen.reset_display()
        self._input_queue.append("look")
        return self.run_until_input()

    def leave_filename_prompt(self) -> str:
        """Return to normal play after a failed or abandoned save/restore filename."""
        self._input_queue.clear()
        self._char_remainder = ""
        self._rewind_pc = None
        if self._restore_undo():
            self.screen.reset_display()
            self._input_queue.append("look")
            return self.run_until_input()
        self._input_queue.append("")
        return self.run_until_input()

    def run_until_input(self) -> str:
        """Advance the VM until input is required or the game exits."""
        self._pending_input_kind = None
        while not self._finished:
            try:
                self._safe_step()
            except NeedInputError as exc:
                self._pending_input_kind = exc.kind
                return self.screen.drain_output()
            except SystemExit:
                self._finished = True
                return self.screen.drain_output()
        return self.screen.drain_output()


class CapturingScreen(vterm.Screen):
    """xyppy Screen that captures text instead of drawing to a terminal."""

    def __init__(
        self,
        env: zenv.Env,
        line_input: Callable[[str], str],
        char_input: Callable[[], str],
    ) -> None:
        super().__init__(env)
        self._line_input = line_input
        self._char_input = char_input
        self._chunks: list[str] = []
        self._last_status: dict[str, str] = {}

    @property
    def last_status(self) -> dict[str, str]:
        return dict(self._last_status)

    def first_draw(self) -> None:
        """Skip terminal setup — we are not using a real console."""

    def reset_display(self) -> None:
        """Clear screen buffers after undo/restore so status text stays in sync."""
        self._chunks.clear()
        self._last_status = {}
        self.textBuf = self.make_screen_buf()
        self.seenBuf = {line: True for line in self.textBuf}
        self.wrapBuf = []
        self.haveNotScrolled = True

    def slow_scroll_effect(self) -> None:
        pass

    def pause_scroll_for_user_input(self) -> None:
        self.flush()
        self.update_seen_lines()

    def overwrite_line_with(self, new_line) -> None:
        pass

    def flush(self) -> None:
        self.finish_wrapping()

    def write(self, text: str) -> None:
        """Capture raw game text before the screen buffer wraps/pads it."""
        if self.env.current_window == 0:
            self._chunks.extend(text)
        super().write(text)

    def write_unwrapped(self, text_as_screenchars, already_seen=False) -> None:
        super().write_unwrapped(text_as_screenchars, already_seen=already_seen)

    def new_line(self) -> None:
        super().new_line()
        if self.env.current_window != 0:
            return
        if not self._chunks or self._chunks[-1] != "\n":
            self._chunks.append("\n")

    def drain_output(self) -> str:
        self.flush()
        text = _clean_zmachine_output("".join(self._chunks))
        self._chunks.clear()
        self._last_status = _status_from_screen(self.env, self.textBuf)
        return text

    def scroll(self, count_lines: bool = True) -> None:
        env = self.env
        self.textBuf.pop(env.top_window_height)
        new_line = self.make_screen_line()
        self.textBuf.append(new_line)
        self.seenBuf[new_line] = False
        self.haveNotScrolled = False

    def scroll_top_line_only(self) -> None:
        env = self.env
        if self.haveNotScrolled and vterm.line_empty(self.textBuf[env.top_window_height]):
            return
        new_line = self.make_screen_line()
        self.textBuf[env.top_window_height] = new_line
        self.seenBuf[new_line] = False
        self.haveNotScrolled = False

    def get_line_of_input(self, prompt: str = "", prefilled: str = "") -> str:
        if prompt:
            self.write(prompt)
            self.flush()
        return self._line_input(prefilled)

    def getch_or_esc_seq(self) -> str:
        self.flush()
        try:
            ch = self._char_input()
            if ch == "\x7f":
                ch = "\b"
            return ch
        except NeedInputError as exc:
            exc.kind = "char"
            raise

    def msg(self, text: str) -> None:
        self.write(text)
        self.flush()
        try:
            self._char_input()
        except NeedInputError:
            raise


def _clean_zmachine_output(text: str) -> str:
    """Trim Z-machine line padding; keep paragraph breaks."""
    lines: list[str] = []
    for line in text.split("\n"):
        stripped = line.strip()
        if stripped == ">":
            continue
        if not stripped:
            if lines and lines[-1] != "":
                lines.append("")
            continue
        lines.append(stripped)

    deduped: list[str] = []
    for line in lines:
        if not deduped or deduped[-1] != line:
            deduped.append(line)

    while deduped and deduped[-1] == "":
        deduped.pop()

    return "\n".join(deduped).strip()


def _status_from_screen(env: zenv.Env, text_buf) -> dict[str, str]:
    """Read the fixed-width Inform status line from the top screen row."""
    if env.top_window_height < 1:
        return {}
    row = text_buf[0]
    line = "".join(c.char for c in row)
    return _parse_status_line(line)


def _parse_status_line(raw: str) -> dict[str, str]:
    """Parse Inform status line: location left, score/moves right."""
    line = raw.rstrip()
    if not line.strip():
        return {}

    score_m = re.search(r"Score:\s*(-?\d+)", line, re.I)
    moves_m = re.search(r"Moves:\s*(\d+)", line, re.I)
    if score_m:
        return {
            "location": line[: score_m.start()].strip(),
            "score": score_m.group(1),
            **({"moves": moves_m.group(1)} if moves_m else {}),
        }
    if moves_m:
        return {
            "location": line[: moves_m.start()].strip(),
            "moves": moves_m.group(1),
        }

    # Compact Inform layout: "Location            <score>/<moves>"
    compact_m = re.search(r"(-?\d+)\s*/\s*(\d+)\s*$", line)
    if compact_m:
        return {
            "location": line[: compact_m.start()].strip(),
            "score": compact_m.group(1),
            "moves": compact_m.group(2),
        }

    collapsed = re.sub(r"\s+", " ", line).strip()
    return {"location": collapsed}