sts2-bot/sts2.py

#!/usr/bin/env python3
"""
sts2.py -- thin client for the STS2MCP localhost interface.

The mod exposes the game on http://localhost:15526 with no auth.

  GET  /api/v1/singleplayer?format=json|markdown   read state
  POST /api/v1/singleplayer {"action": ..., ...}   perform one action

Learning from the state model: every response carries `state_type`, and
each state_type has a small fixed action set. That set IS our action
space, so this module also maps state_type -> legal actions.

The loop is strictly closed. Playing a card removes it from hand and
shifts every later index, so we observe -> act once -> observe again.
Never precompute an action list.
"""

from __future__ import annotations

import json
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any

BASE = "http://localhost:15526"
TIMEOUT = 20.0


class Sts2Error(RuntimeError):
    pass


# state_type -> legal actions. Source: STS2MCP docs/raw-simplified.md
LEGAL_ACTIONS: dict[str, tuple[str, ...]] = {
    "menu": ("menu_select",),
    "game_over": ("menu_select",),
    "monster": ("play_card", "use_potion", "end_turn"),
    "elite": ("play_card", "use_potion", "end_turn"),
    "boss": ("play_card", "use_potion", "end_turn"),
    "hand_select": ("combat_select_card", "combat_confirm_selection"),
    "rewards": ("claim_reward", "proceed"),
    "card_reward": ("select_card_reward", "skip_card_reward"),
    "map": ("choose_map_node",),
    "event": ("choose_event_option", "advance_dialogue"),
    "rest_site": ("choose_rest_option", "proceed"),
    "shop": ("shop_purchase", "proceed"),
    "fake_merchant": ("shop_purchase", "proceed"),
    "treasure": ("claim_treasure_relic", "proceed"),
    "card_select": ("select_card", "confirm_selection", "cancel_selection"),
    "bundle_select": (
        "select_bundle",
        "confirm_bundle_selection",
        "cancel_bundle_selection",
    ),
    "relic_select": ("select_relic", "skip_relic_selection"),
    "crystal_sphere": (
        "crystal_sphere_set_tool",
        "crystal_sphere_click_cell",
        "crystal_sphere_proceed",
    ),
    "overlay": (),
    "unknown": (),
}

COMBAT_STATES = ("monster", "elite", "boss")


@dataclass
class ActionResult:
    status: str
    message: str
    raw: dict

    @property
    def ok(self) -> bool:
        return self.status == "ok"


def _request(path: str, payload: dict | None = None, timeout: float = TIMEOUT) -> dict:
    url = BASE + path
    if payload is None:
        request = urllib.request.Request(url, method="GET")
    else:
        request = urllib.request.Request(
            url,
            data=json.dumps(payload).encode("utf-8"),
            method="POST",
            headers={"Content-Type": "application/json"},
        )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = response.read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")
        raise Sts2Error(f"HTTP {exc.code} on {path}: {detail[:300]}") from exc
    except urllib.error.URLError as exc:
        raise Sts2Error(
            f"cannot reach the game on {BASE}. Is STS2 running with the mod loaded? ({exc})"
        ) from exc

    try:
        return json.loads(body)
    except json.JSONDecodeError as exc:
        raise Sts2Error(f"non-JSON response from {path}: {body[:200]!r}") from exc


# --------------------------------------------------------------------------
# Reads
# --------------------------------------------------------------------------

def is_up() -> bool:
    try:
        data = _request("/", timeout=5.0)
        return data.get("status") == "ok"
    except Sts2Error:
        return False


def state(fmt: str = "json") -> dict:
    """Current game state. fmt: 'json' or 'markdown' (markdown returns text)."""
    if fmt == "markdown":
        req = urllib.request.Request(f"{BASE}/api/v1/singleplayer?format=markdown")
        with urllib.request.urlopen(req, timeout=TIMEOUT) as response:
            return {"markdown": response.read().decode("utf-8")}
    return _request("/api/v1/singleplayer?format=json")


def state_type(observation: dict) -> str:
    return observation.get("state_type", "unknown")


def legal_actions(observation: dict) -> tuple[str, ...]:
    return LEGAL_ACTIONS.get(state_type(observation), ())


def profile() -> dict:
    return _request("/api/v1/profile")


def compendium() -> dict:
    return _request("/api/v1/compendium")


def wiki(query: str, item_type: str = "all", limit: int = 10) -> dict:
    from urllib.parse import quote

    return _request(
        f"/api/v1/wiki?query={quote(query)}&item_type={item_type}&limit={limit}"
    )


def profiles() -> dict:
    return _request("/api/v1/profiles")


# --------------------------------------------------------------------------
# Writes -- one action at a time
# --------------------------------------------------------------------------

def act(action: str, **params: Any) -> ActionResult:
    payload = {"action": action}
    payload.update(params)
    data = _request("/api/v1/singleplayer", payload)
    return ActionResult(
        status=data.get("status", "?"),
        # Errors arrive under `error`, not `message`. Reading only `message`
        # makes every rejection look empty.
        message=data.get("message") or data.get("error") or "",
        raw=data,
    )


def play_card(card_index: int, target: str | None = None) -> ActionResult:
    params: dict[str, Any] = {"card_index": card_index}
    if target is not None:
        params["target"] = target
    return act("play_card", **params)


def end_turn() -> ActionResult:
    return act("end_turn")


def use_potion(slot: int, target: str | None = None) -> ActionResult:
    params: dict[str, Any] = {"slot": slot}
    if target is not None:
        params["target"] = target
    return act("use_potion", **params)


def menu_select(option: str, seed: str | None = None) -> ActionResult:
    params: dict[str, Any] = {"option": option}
    if seed is not None:
        params["seed"] = seed
    return act("menu_select", **params)


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def _main() -> int:
    import sys

    if not is_up():
        print(f"game not reachable at {BASE}")
        return 1

    if len(sys.argv) > 1 and sys.argv[1] == "--dump":
        obs = state()
        print(json.dumps(obs, indent=2))
        return 0

    obs = state()
    st = state_type(obs)
    print(f"state_type : {st}")
    print(f"run        : {obs.get('run')}")
    print(f"legal      : {', '.join(legal_actions(obs)) or '(none - manual)'}")
    if "options" in obs:
        print(f"options    : {obs['options']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(_main())