game.py

"""Doodle Duel game logic — human draws, AI guesses via vision model."""
from __future__ import annotations

import random
import re
import time
from dataclasses import dataclass, field

import numpy as np
from PIL import Image

import config
import vision_client


@dataclass
class DoodleState:
    status: str = "idle"           # idle | choosing | drawing | won | lost
    difficulty: str = "medium"     # easy | medium | hard
    word: str = ""
    emoji: str = ""
    category: str = ""
    accepted: list = field(default_factory=list)
    options: list = field(default_factory=list)
    hints_used: int = 0
    revealed: set = field(default_factory=set)
    tried: list = field(default_factory=list)
    deck: list = field(default_factory=list)
    guess_log: list = field(default_factory=list)
    reason_log: list = field(default_factory=list)
    last_caught: str = ""
    stream_think: str = ""
    last_canvas: object = None     # latest canvas captured on stroke-end

    start: float = 0.0
    time_left: float = config.ROUND_SECONDS
    last_guess_t: float = 0.0
    guess_in_flight: bool = False
    round: int = 0
    score: int = 0
    best: int = 0
    streak: int = 0
    best_streak: int = 0


_ARTICLES = {"a", "an", "the"}


def _norm(s: str) -> str:
    s = re.sub(r"[^a-z ]", "", s.lower())
    return " ".join(t for t in s.split() if t not in _ARTICLES)


def matches(accepted, guesses):
    acc = {_norm(a) for a in accepted}
    for g in guesses:
        if _norm(g) in acc:
            return g
    return ""


def to_pil(canvas):
    if canvas is None:
        return None
    if not isinstance(canvas, dict):
        if isinstance(canvas, np.ndarray):
            return Image.fromarray(canvas.astype("uint8"))
        return canvas
    # Prefer composite (background + all drawn layers merged)
    comp = canvas.get("composite")
    if comp is None:
        comp = canvas.get("background")
    if comp is None:
        return None
    pil = Image.fromarray(comp.astype("uint8")) if isinstance(comp, np.ndarray) else comp
    # If composite appears blank, check whether any drawing layer has strokes
    if is_blank(pil):
        for layer in (canvas.get("layers") or []):
            if layer is None:
                continue
            layer_arr = layer if isinstance(layer, np.ndarray) else np.asarray(layer)
            if layer_arr.ndim == 3 and layer_arr.shape[2] == 4:
                if (layer_arr[:, :, 3] > 10).mean() > 0.001:
                    # Compose layer onto white background for the model
                    bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
                    merged = Image.alpha_composite(bg, Image.fromarray(layer_arr.astype("uint8")))
                    return merged.convert("RGB")
    return pil


def is_blank(pil) -> bool:
    if pil is None:
        return True
    arr = np.asarray(pil.convert("L"))
    return (arr < 250).mean() < 0.002


def letter_idxs(word):
    return [i for i, ch in enumerate(word) if ch not in set(" -")]


def masked(word, revealed) -> str:
    return "".join(ch if ch in set(" -") else (ch if i in revealed else "-")
                   for i, ch in enumerate(word))


def fits_pattern(guess, word, revealed) -> bool:
    wl = [(i, ch.lower()) for i, ch in enumerate(word) if ch not in set(" -")]
    gl = [ch.lower() for ch in guess if ch.isalpha()]
    if len(gl) != len(wl) or len(guess.split()) != len(word.split()):
        return False
    for pos, (orig_i, ch) in enumerate(wl):
        if orig_i in revealed and gl[pos] != ch:
            return False
    return True


def pattern_text(state: DoodleState) -> str:
    n = len(letter_idxs(state.word))
    pat = masked(state.word, state.revealed).replace(" ", " / ")
    return (f" The answer has {n} letters and fits this pattern (each dash is one "
            f"unknown letter, / separates words): {pat}")


def _word_pool(difficulty: str) -> list:
    if difficulty == "easy":
        return [w for w, (_, _, _, d) in config.WORDS.items() if d == "easy"]
    if difficulty == "hard":
        return [w for w, (_, _, _, d) in config.WORDS.items() if d == "hard"]
    return list(config.WORDS.keys())  # medium = all


def _deal(state: DoodleState) -> list:
    if len(state.deck) < config.WORD_CHOICES:
        state.deck = _word_pool(state.difficulty)
        random.shuffle(state.deck)
    words = state.deck[:config.WORD_CHOICES]
    state.deck = state.deck[config.WORD_CHOICES:]
    return words


def offer_words(state: DoodleState, difficulty: str = "medium") -> DoodleState:
    if state.difficulty != difficulty:
        state.deck = []          # rebuild deck when switching difficulty
    state.difficulty = difficulty
    state.options = _deal(state)
    state.status = "choosing"
    state.guess_log = []
    state.reason_log = []
    state.last_caught = ""
    state.stream_think = ""
    return state


def choose_word(state: DoodleState, idx: int) -> DoodleState:
    if state.status != "choosing" or not (0 <= idx < len(state.options)):
        return state
    word = state.options[idx]
    emoji, category, accepted, _ = config.WORDS[word]
    state.status = "drawing"
    state.word, state.emoji, state.category, state.accepted = word, emoji, category, accepted
    state.start = time.monotonic()
    state.time_left = config.ROUND_SECONDS
    state.last_guess_t = 0.0
    state.hints_used = 0
    state.revealed = set()
    state.tried = []
    state.guess_log = []
    state.reason_log = []
    state.last_caught = ""
    state.stream_think = ""
    state.guess_in_flight = False
    state.last_canvas = None
    state.round += 1
    return state


def use_hint(state: DoodleState) -> DoodleState:
    if state.status == "drawing" and state.hints_used < config.MAX_HINTS:
        idxs = letter_idxs(state.word)
        hidden = [i for i in idxs if i not in state.revealed]
        if hidden:
            k = max(1, round(len(idxs) * 0.25))
            for i in random.sample(hidden, min(k, len(hidden))):
                state.revealed.add(i)
        state.hints_used += 1
    return state


def _do_guess(state: DoodleState, pil):
    guesses, reason = vision_client.guess(pil, hint=pattern_text(state), avoid=state.tried)
    if reason:
        state.reason_log = (state.reason_log + [reason])[-6:]
    if not guesses:
        return
    for g in guesses:
        gl = g.lower().strip()
        if gl and gl not in state.tried:
            state.tried.append(gl)
    hit = matches(state.accepted, guesses)
    if hit:
        base = max(10, int(state.time_left * 5) + 20)
        gained = max(5, int(base * (1 - 0.3 * state.hints_used)))
        state.score += gained
        state.best = max(state.best, state.score)
        state.streak += 1
        state.best_streak = max(state.best_streak, state.streak)
        state.last_caught = f"{hit} (+{gained})"
        state.status = "won"
        return
    # Show all guesses as chips so the user always sees the robot's attempts
    state.guess_log = (state.guess_log + guesses)[-8:]


def poll_guess(state: DoodleState, canvas) -> DoodleState:
    """One guess from `canvas`, updating reason/guess logs and win state.
    No internal time-gate — the background loop controls the cadence. Skips
    blank canvases so we never waste a call (or tokens) on an empty board."""
    if state.status != "drawing":
        return state
    state.time_left = max(0.0, config.ROUND_SECONDS - (time.monotonic() - state.start))
    if state.time_left <= 0:
        state.status = "lost"
        state.streak = 0
        return state
    pil = to_pil(canvas)
    if pil is None or is_blank(pil):
        return state
    _do_guess(state, pil)
    return state


def stash_canvas(state: DoodleState, canvas) -> DoodleState:
    """Called on canvas.change (stroke-end). Cheap: just stores the latest
    canvas. NO network — drawing must never trigger a model request (the
    per-stroke calls were what flooded the Space and tripped HF's 429). The
    timer is the only thing that guesses, reading from state.last_canvas."""
    if state.status != "drawing":
        return state
    state.last_canvas = canvas
    return state


def tick_time(state: DoodleState):
    """Timer-only tick — just updates time_left, no guessing."""
    if state.status != "drawing":
        return state
    state.time_left = max(0.0, config.ROUND_SECONDS - (time.monotonic() - state.start))
    if state.time_left <= 0:
        state.status = "lost"
        state.streak = 0
    return state