Bolt-Pet/bolt_pet/ui/sprite.py

"""Sprite loading + frame animation.

Convention: assets/sprites/<state>/*.png, frames played in filename-sorted
order (e.g. frame_00.png, frame_01.png, ...), looping. <state> matches
bolt_pet.state.PetState values: idle, listening, thinking, talking.

If a state's directory has no frames (real art not dropped in yet), falls
back to a small procedurally-drawn placeholder blob so the app still runs
end-to-end. Swap in real sprite sheets by pointing SPRITE_DIR at your own
folder (see assets/sprites/README.md) — no code changes needed as long as
the same per-state-subfolder-of-PNGs convention is followed. If your sheets
use a different layout (single grid image, etc.), tell me the format and
this loader can be adapted.
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional

from PySide6.QtCore import QSize, Qt
from PySide6.QtGui import QColor, QPainter, QPixmap

from ..state import PetState

DEFAULT_SPRITE_DIR = Path(__file__).resolve().parent.parent / "assets" / "sprites"

# Placeholder palette per state, used only when no frames are found.
_PLACEHOLDER_COLORS = {
    PetState.IDLE: QColor(120, 170, 240),
    PetState.LISTENING: QColor(120, 220, 160),
    PetState.THINKING: QColor(230, 190, 90),
    PetState.TALKING: QColor(240, 130, 150),
    PetState.ERROR: QColor(220, 90, 90),
}


def _placeholder_frames(state: PetState, size: int) -> list[QPixmap]:
    """A tiny 2-frame "breathing" blob so idle/listening/etc. are visually
    distinguishable even before real art exists."""
    color = _PLACEHOLDER_COLORS.get(state, QColor(150, 150, 150))
    frames = []
    for scale in (1.0, 0.92):
        pixmap = QPixmap(size, size)
        pixmap.fill(Qt.transparent)
        painter = QPainter(pixmap)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.setBrush(color)
        painter.setPen(Qt.NoPen)
        margin = size * (1 - scale) / 2
        painter.drawEllipse(int(margin), int(margin), int(size * scale), int(size * scale))
        # simple eyes so it reads as a face, not just a circle
        eye_r = max(2, size // 16)
        eye_y = int(size * 0.42)
        painter.setBrush(QColor(30, 30, 40))
        painter.drawEllipse(int(size * 0.36) - eye_r, eye_y - eye_r, eye_r * 2, eye_r * 2)
        painter.drawEllipse(int(size * 0.64) - eye_r, eye_y - eye_r, eye_r * 2, eye_r * 2)
        painter.end()
        frames.append(pixmap)
    return frames


class SpriteAnimation:
    """One state's frame sequence + current playback position."""

    def __init__(self, frames: list[QPixmap]):
        self.frames = frames or []
        self._index = 0

    def advance(self) -> None:
        if self.frames:
            self._index = (self._index + 1) % len(self.frames)

    def current(self) -> Optional[QPixmap]:
        if not self.frames:
            return None
        return self.frames[self._index]

    def reset(self) -> None:
        self._index = 0


def _load_frames_from_dir(directory: Path, size: int) -> list[QPixmap]:
    if not directory.is_dir():
        return []
    paths = sorted(directory.glob("*.png")) + sorted(directory.glob("*.PNG"))
    frames = []
    for path in paths:
        pixmap = QPixmap(str(path))
        if pixmap.isNull():
            continue
        if pixmap.size() != QSize(size, size):
            pixmap = pixmap.scaled(size, size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        frames.append(pixmap)
    return frames


WALK = "walk"

# Animations that aren't pipeline states. Walking is a property of *movement*,
# orthogonal to whether the pet is idle/listening/talking, so it deliberately
# isn't a PetState — state.py stays a description of the conversation, not of
# the body. Loaded the same way, keyed by name.
EXTRA_ANIMATIONS = (WALK,)


class SpriteSet:
    """All animations for every PetState, plus the extras, from *sprite_dir*."""

    def __init__(self, sprite_dir: Path = DEFAULT_SPRITE_DIR, size: int = 160):
        self.size = size
        self._animations: dict[str, SpriteAnimation] = {}
        self._loaded: set[str] = set()  # keys backed by real art, not placeholders
        for state in PetState:
            frames = _load_frames_from_dir(sprite_dir / state.value, size)
            if frames:
                self._loaded.add(state.value)
            else:
                frames = _placeholder_frames(state, size)
            self._animations[state.value] = SpriteAnimation(frames)
        for name in EXTRA_ANIMATIONS:
            frames = _load_frames_from_dir(sprite_dir / name, size)
            if frames:
                self._loaded.add(name)
                self._animations[name] = SpriteAnimation(frames)

    @staticmethod
    def _key(key) -> str:
        return key.value if isinstance(key, PetState) else str(key)

    def get(self, key) -> SpriteAnimation:
        """Animation for a PetState or an extra name. Unknown/absent extras
        fall back to idle, so a sprite folder with no walk/ still runs."""
        return self._animations.get(self._key(key)) or self._animations[PetState.IDLE.value]

    def has(self, key) -> bool:
        """True only when real frames were found — the caller uses this to
        decide whether to use an extra animation at all, rather than being
        handed a placeholder blob that looks nothing like walking."""
        return self._key(key) in self._loaded