vision_client.py

"""Vision client — asks the model to guess what a doodle depicts,
and draws SVG pictures for the player to guess.

guess(pil_image) -> list[str] of lowercase one/two-word guesses (best first).
draw_svg(word) -> str of SVG markup (or "" on failure).
"""
from __future__ import annotations

import base64
import io
import random
import re
import sys

import requests

import config

_resolved_model = None      # cached result of model-name auto-detection
_warned = set()             # so we only print each backend warning once


def _warn_once(msg: str):
    if msg not in _warned:
        _warned.add(msg)
        print(f"[vision_client] {msg}", file=sys.stderr)


def _auth_headers():
    h = {"Content-Type": "application/json"}
    if config.MODEL_API_KEY:
        h["Authorization"] = f"Bearer {config.MODEL_API_KEY}"
    return h


def _model_name():
    """The id to send in the `model` field. If MODEL_NAME is 'auto' (or empty),
    ask the server's /v1/models and use what it actually serves — llama-server
    ignores this field, but vLLM validates it, so auto-detection lets one config
    drive either backend (and the llama.cpp gguf name need not be known)."""
    global _resolved_model
    name = (config.MODEL_NAME or "").strip()
    if name and name.lower() != "auto":
        return name
    if _resolved_model:
        return _resolved_model
    try:
        r = requests.get(f"{config.MODEL_BASE_URL}/v1/models",
                         headers=_auth_headers(), timeout=10)
        r.raise_for_status()
        _resolved_model = r.json()["data"][0]["id"]
    except Exception:
        _resolved_model = "default"      # llama-server accepts any id
    return _resolved_model


def health():
    """(ok, detail) for the configured backend — used for a startup check."""
    if config.MOCK_MODE:
        return True, "MOCK mode (no model server set)"
    base = config.MODEL_BASE_URL
    for path in ("/health", "/v1/models"):    # llama.cpp has /health; vLLM has /v1/models
        try:
            r = requests.get(f"{base}{path}", headers=_auth_headers(), timeout=5)
            if r.ok:
                return True, f"{base} — model '{_model_name()}'"
        except Exception as e:
            last = e.__class__.__name__
    return False, f"{base} unreachable ({locals().get('last', 'no response')})"

_GUESS_INTRO = (
    "We are playing Pictionary. This is a rough, simple doodle drawn by a human. "
    "Guess what or WHO it depicts — it may be an object, a brand/logo, a character, "
    "or a famous person (use clues like jersey numbers, initials, symbols). Be "
    "SPECIFIC: if it is a soccer ball say 'soccer ball' not 'ball'; if the clues "
    "point to a person, name them."
)

_GUESS_FORMAT = (
    "You are told the answer's letter count, number of words, and a dash pattern "
    "(revealed letters are filled in) — every guess MUST have exactly that many "
    "letters and words and match the revealed letters.\n"
    "Respond in EXACTLY two lines, nothing else:\n"
    "THINK: one short sentence describing the shapes you see and your reasoning — REQUIRED, always include this even if uncertain.\n"
    "GUESS: your top 3 guesses, most likely first, comma-separated, lowercase."
)

_DRAW_SVG_PROMPT = """Draw "{word}" as a Pictionary SVG that players can CLEARLY recognize and guess.

Think step by step:
1. What is the most iconic silhouette or shape of "{word}"?
2. What 2-3 key details make it instantly recognizable?
3. What realistic colors should I use?

Then draw it.

SVG rules:
- viewBox="0 0 500 360" width="500" height="360"
- The canvas background is already WHITE — do NOT add a background rect and do NOT set fill on the <svg> element
- Allowed child elements: circle, rect, path, line, ellipse, polygon, polyline
- 8-35 elements — enough detail to be recognizable
- NO text, letters, numbers, or labels of any kind
- Use realistic, VISIBLE colors on the white background (e.g. red body, black spots for ladybug)
- Use stroke-width="2" or "3" for clean outlines on white
- Center the subject, leave a small margin

The goal is for a human to look at your drawing and immediately know it is "{word}".
Draw it like a clear, clean illustration — NOT abstract, NOT dark.

Output ONLY the raw SVG. First character must be < and last must be >."""


def _to_data_url(pil_image) -> str:
    buf = io.BytesIO()
    pil_image.convert("RGB").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"


def _split_guesses(text: str):
    text = re.sub(r"[\n;]", ",", text.lower())
    out = []
    for part in text.split(","):
        p = re.sub(r"[^a-z ]", "", part).strip()
        if p and p not in out:
            out.append(p)
    return out[:3]


def _parse_response(text: str):
    """Pull (guesses, reasoning) from the model's two-line THINK/GUESS reply.
    Falls back gracefully if the model ignores the format."""
    # Drop native CoT blocks (TGI and llama.cpp both use these)
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    # Qwen3 on llama.cpp sometimes emits "Thinking Process:\n...\n\n" before the answer
    text = re.sub(r"(?i)thinking process:.*?\n\n", "", text, flags=re.DOTALL)
    reason, guess_line = "", ""
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        low = line.lower()
        if low.startswith("think"):
            reason = line.split(":", 1)[1].strip() if ":" in line else line
        elif low.startswith("guess"):
            guess_line = line.split(":", 1)[1].strip() if ":" in line else line
    if not guess_line:
        lines = [l.strip() for l in text.splitlines() if l.strip()]
        if lines:
            guess_line = lines[-1]
            if not reason and len(lines) > 1:
                reason = " ".join(lines[:-1])
    return _split_guesses(guess_line), reason.strip()


def guess(pil_image, hint="", avoid=None):
    """Returns (guesses: list[str], reasoning: str)."""
    if config.MOCK_MODE or pil_image is None:
        return random.sample(list(config.WORDS.keys()), 3), "(mock) just guessing!"
    url = f"{config.MODEL_BASE_URL}/v1/chat/completions"
    headers = _auth_headers()
    avoid_txt = ""
    if avoid:
        avoid_txt = (
            " ALREADY GUESSED AND WRONG — do NOT use any of these words, not even "
            "as part of a different phrase: " + ", ".join(avoid) + "."
            " You MUST guess something completely different."
        )
    prompt = _GUESS_INTRO + avoid_txt + (hint or "") + "\n" + _GUESS_FORMAT
    payload = {
        "model": _model_name(),
        "messages": [
            {"role": "system", "content": "/no_think\nYou are a Pictionary AI. Be concise."},
            {"role": "user", "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": _to_data_url(pil_image)}},
            ]},
        ],
        "max_tokens": 200,
        "temperature": 0.6,
        "chat_template_kwargs": {"enable_thinking": False},
    }
    for attempt in range(2):
        try:
            r = requests.post(url, json=payload, headers=headers, timeout=config.REQUEST_TIMEOUT)
            r.raise_for_status()
            msg = r.json()["choices"][0]["message"]
            txt = msg.get("content") or msg.get("reasoning_content") or ""
            return _parse_response(txt)
        except requests.RequestException as e:
            if attempt == 0:
                continue
            _warn_once(f"guess request failed against {config.MODEL_BASE_URL}: {e}")
            return [], ""
    return [], ""


# ── Mock SVG shapes for testing without the real model ───────────────────────
_MOCK_SVGS = {
    "cat": '<svg viewBox="0 0 500 360" width="500" height="360" xmlns="http://www.w3.org/2000/svg"><ellipse cx="250" cy="220" rx="100" ry="80" fill="#f4a460" stroke="#8b4513" stroke-width="3"/><circle cx="250" cy="150" r="70" fill="#f4a460" stroke="#8b4513" stroke-width="3"/><polygon points="195,95 175,45 220,85" fill="#f4a460" stroke="#8b4513" stroke-width="2"/><polygon points="305,95 325,45 280,85" fill="#f4a460" stroke="#8b4513" stroke-width="2"/><circle cx="220" cy="145" r="12" fill="#333"/><circle cx="280" cy="145" r="12" fill="#333"/><ellipse cx="250" cy="175" rx="18" ry="12" fill="#ff9999"/><line x1="210" y1="175" x2="160" y2="165" stroke="#8b4513" stroke-width="2"/><line x1="210" y1="180" x2="155" y2="182" stroke="#8b4513" stroke-width="2"/><line x1="290" y1="175" x2="340" y2="165" stroke="#8b4513" stroke-width="2"/><line x1="290" y1="180" x2="345" y2="182" stroke="#8b4513" stroke-width="2"/></svg>',
    "sun": '<svg viewBox="0 0 500 360" width="500" height="360" xmlns="http://www.w3.org/2000/svg"><circle cx="250" cy="180" r="70" fill="#FFD700" stroke="#FFA500" stroke-width="4"/><line x1="250" y1="80" x2="250" y2="50" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="250" y1="280" x2="250" y2="310" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="150" y1="180" x2="120" y2="180" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="350" y1="180" x2="380" y2="180" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="179" y1="109" x2="158" y2="88" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="321" y1="109" x2="342" y2="88" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="179" y1="251" x2="158" y2="272" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/><line x1="321" y1="251" x2="342" y2="272" stroke="#FFA500" stroke-width="5" stroke-linecap="round"/></svg>',
    "house": '<svg viewBox="0 0 500 360" width="500" height="360" xmlns="http://www.w3.org/2000/svg"><rect x="120" y="200" width="260" height="150" fill="#DEB887" stroke="#8B4513" stroke-width="3"/><polygon points="100,200 250,80 400,200" fill="#CD5C5C" stroke="#8B0000" stroke-width="3"/><rect x="200" y="260" width="100" height="90" fill="#8B4513"/><rect x="140" y="230" width="60" height="55" fill="#87CEEB" stroke="#8B4513" stroke-width="2"/><rect x="300" y="230" width="60" height="55" fill="#87CEEB" stroke="#8B4513" stroke-width="2"/></svg>',
}

_SVG_EXTRACT = re.compile(r"(<svg[\s\S]*?</svg>)", re.IGNORECASE)
_DARK_FILL = re.compile(
    r'fill="(?:#0{3,6}|black|#1[0-9a-f]{5}|#2[0-9a-f]{5}|rgb\(0,\s*0,\s*0\))"',
    re.IGNORECASE,
)


def _sanitize_svg(svg: str) -> str:
    """Remove dark background fills that would hide the drawing."""
    # Strip fill from the <svg> root tag (make it transparent)
    svg = re.sub(r'(<svg\b[^>]*?)\s+fill="[^"]*"', r'\1', svg, count=1, flags=re.IGNORECASE)
    # Remove any full-canvas background rect with a dark fill
    def _maybe_drop_rect(m):
        s = m.group(0)
        is_fullsize = (
            re.search(r'width="(?:500|100%)"', s) and
            re.search(r'height="(?:360|100%)"', s)
        )
        return '' if (is_fullsize and _DARK_FILL.search(s)) else s
    svg = re.sub(r'<rect\b[^>]*/>', _maybe_drop_rect, svg, flags=re.IGNORECASE)
    return svg


def draw_svg(word: str) -> str:
    """Ask the model to generate an SVG drawing of `word`.
    Returns raw SVG markup, or a mock SVG in MOCK_MODE."""
    if config.MOCK_MODE:
        w = word.lower()
        if w in _MOCK_SVGS:
            return _MOCK_SVGS[w]
        return _MOCK_SVGS.get("house")   # fallback mock

    url = f"{config.MODEL_BASE_URL}/v1/chat/completions"
    prompt = _DRAW_SVG_PROMPT.format(word=word)
    payload = {
        "model": _model_name(),
        "messages": [
            {"role": "system", "content": "You are an SVG illustrator making clear Pictionary drawings."},
            {"role": "user", "content": prompt},
        ],
        "max_tokens": 6000,   # thinking + SVG needs room
        "temperature": 0.7,
    }
    try:
        r = requests.post(url, json=payload, headers=_auth_headers(),
                          timeout=config.REQUEST_TIMEOUT)
        r.raise_for_status()
        msg = r.json()["choices"][0]["message"]
        txt = msg.get("content") or msg.get("reasoning_content") or ""
        # Search the FULL response first — model sometimes puts SVG inside <think>
        m = _SVG_EXTRACT.search(txt)
        if not m:
            # Fall back: strip thinking blocks and try again
            txt2 = re.sub(r"<think>.*?</think>", "", txt, flags=re.DOTALL)
            txt2 = re.sub(r"(?i)thinking process:.*?\n\n", "", txt2, flags=re.DOTALL)
            m = _SVG_EXTRACT.search(txt2)
        if m:
            svg = m.group(1)
            # Normalize viewport so it fits our display area
            svg = re.sub(r'width="[^"]*"', 'width="500"', svg)
            svg = re.sub(r'height="[^"]*"', 'height="360"', svg)
            if 'viewBox' not in svg:
                svg = svg.replace('<svg', '<svg viewBox="0 0 500 360"', 1)
            svg = _sanitize_svg(svg)
            return svg
    except Exception as e:
        _warn_once(f"draw_svg request failed: {e}")
    return ""