app.py

import json
import io
import base64
import random
import re
import os
import time
import logging
import urllib.parse
import urllib.request

import gradio as gr
from huggingface_hub import InferenceClient

try:
    import spaces
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(fn=None, **_kwargs):
            if callable(fn):
                return fn

            def _decorator(func):
                return func

            return _decorator

    spaces = _SpacesShim()


PYXEL_SCRIPT = """
import json
import js
import pyxel

W = 16
H = 16
PIX = 12
CANVAS_X = 8
CANVAS_Y = 8
PALETTE_X = CANVAS_X + W * PIX + 16
PALETTE_Y = 8
PALETTE_COLS = 3
PALETTE_ROWS = 8
SWATCH_W = 14
SWATCH_H = 12
SWATCH_GAP_X = 4
SWATCH_GAP_Y = 4
TOOL_Y = PALETTE_Y + PALETTE_ROWS * (SWATCH_H + SWATCH_GAP_Y) + 8

PALETTE_HEX = [
    0xF7F2E8,  # 00 paper white
    0x2E2A26,  # 01 black
    0x8A877F,  # 02 warm gray
    0xCC3B3B,  # 03 red
    0xA72F4C,  # 04 carmine
    0xE1823A,  # 05 orange
    0xEBCF55,  # 06 yellow
    0xA8C857,  # 07 yellow green
    0x63B24D,  # 08 leaf green
    0x2F8F4E,  # 09 emerald green
    0x2B5D3F,  # 10 deep green
    0x36A69A,  # 11 turquoise
    0x6C9EDB,  # 12 sky blue
    0x3F61B8,  # 13 cobalt blue
    0x7A57B5,  # 14 violet
    0x7A4A2B,  # 15 brown
    0x2E3E8C,  # 16 ultramarine
    0xA5479B,  # 17 magenta
    0xE9B28F,  # 18 peach
    0xB89043,  # 19 ochre
    0xD2A47E,  # 20 skin tone
    0x4A2D1F,  # 21 dark brown
    0xCFCBC3,  # 22 light gray
    0x9FA6B2,  # 23 cool gray
]

# Fallback blend pairs for colors 16-23 when runtime supports only 16 colors.
DITHER_FALLBACK = {
    16: (13, 14),
    17: (4, 14),
    18: (0, 5),
    19: (6, 15),
    20: (0, 15),
    21: (1, 15),
    22: (0, 2),
    23: (1, 2),
}

pixels = [0] * (W * H)
selected_color = 7
erase_mode = False
runtime_color_cap = 16
export_lamp_frames = 0
EXPORT_LAMP_MAX_FRAMES = 18
ai_phase = "idle"  # idle | judging | success | fail
ai_anim_frame = 0
ai_result_frames = 0
AI_RESULT_MAX_FRAMES = 90


def setup_palette():
    # Assign 24 custom colors if the runtime supports palette extension.
    # If not, only the supported color slots are updated.
    global runtime_color_cap
    runtime_color_cap = 0
    for i, color in enumerate(PALETTE_HEX):
        try:
            pyxel.colors[i] = color
            runtime_color_cap = i + 1
        except Exception:
            break
    if runtime_color_cap <= 0:
        runtime_color_cap = 16


def to_runtime_color(index):
    if index < 0:
        return 0
    if index < runtime_color_cap:
        return index
    return index % runtime_color_cap


def draw_logical_color_rect(x, y, w, h, logical_index):
    if logical_index < runtime_color_cap:
        pyxel.rect(x, y, w, h, logical_index)
        return

    c1, c2 = DITHER_FALLBACK.get(
        logical_index,
        (logical_index % runtime_color_cap, (logical_index + 1) % runtime_color_cap),
    )
    pyxel.rect(x, y, w, h, c1)
    for dy in range(h):
        for dx in range(w):
            if (dx + dy) % 2 == 1:
                pyxel.pset(x + dx, y + dy, c2)


def mouse_pos_to_cell(mx, my):
    if mx < CANVAS_X or my < CANVAS_Y:
        return None
    cx = (mx - CANVAS_X) // PIX
    cy = (my - CANVAS_Y) // PIX
    if 0 <= cx < W and 0 <= cy < H:
        return int(cx), int(cy)
    return None


def paint_at(mx, my):
    pos = mouse_pos_to_cell(mx, my)
    if pos is None:
        return
    x, y = pos
    pixels[y * W + x] = 0 if erase_mode else selected_color


def export_payload():
    payload = json.dumps(
        {
            "width": W,
            "height": H,
            "pixels": pixels,
            "palette_hex": PALETTE_HEX,
        }
    )
    js.window.__PYXEL_PAYLOAD = payload


def clear_all():
    for i in range(W * H):
        pixels[i] = 0


def sync_ai_phase_from_js():
    global ai_phase, ai_anim_frame, ai_result_frames

    js_phase = ""
    try:
        js_phase = str(js.window.__AI_JUDGE_STATE or "")
    except Exception:
        js_phase = ""

    if js_phase in ("idle", "judging", "success", "fail") and js_phase != ai_phase:
        ai_phase = js_phase
        ai_anim_frame = 0
        if js_phase == "idle":
            ai_result_frames = 0
        if js_phase in ("success", "fail"):
            ai_result_frames = AI_RESULT_MAX_FRAMES


def update():
    global selected_color, erase_mode, export_lamp_frames, ai_anim_frame, ai_phase, ai_result_frames

    sync_ai_phase_from_js()
    ai_anim_frame += 1
    if ai_phase in ("success", "fail") and ai_result_frames > 0:
        ai_result_frames -= 1
        if ai_result_frames <= 0:
            ai_phase = "idle"

    mx = pyxel.mouse_x
    my = pyxel.mouse_y

    if pyxel.btn(pyxel.MOUSE_BUTTON_LEFT):
        paint_at(mx, my)

    if export_lamp_frames > 0:
        export_lamp_frames -= 1

    # 24-color selection area (3x8)
    if pyxel.btnp(pyxel.MOUSE_BUTTON_LEFT):
        for i in range(24):
            px = PALETTE_X + (i % PALETTE_COLS) * (SWATCH_W + SWATCH_GAP_X)
            py = PALETTE_Y + (i // PALETTE_COLS) * (SWATCH_H + SWATCH_GAP_Y)
            if px <= mx < px + SWATCH_W and py <= my < py + SWATCH_H:
                selected_color = i
                erase_mode = False

        # Eraser toggle
        if PALETTE_X <= mx < PALETTE_X + 52 and TOOL_Y <= my < TOOL_Y + 18:
            erase_mode = True

        # Export button
        if PALETTE_X <= mx < PALETTE_X + 52 and TOOL_Y + 24 <= my < TOOL_Y + 42:
            export_lamp_frames = EXPORT_LAMP_MAX_FRAMES
            export_payload()

        # Clear button
        if PALETTE_X <= mx < PALETTE_X + 52 and TOOL_Y + 48 <= my < TOOL_Y + 66:
            clear_all()


def draw():
    pyxel.cls(1)

    # Draw grid
    for y in range(H):
        for x in range(W):
            c = pixels[y * W + x]
            draw_logical_color_rect(
                CANVAS_X + x * PIX,
                CANVAS_Y + y * PIX,
                PIX - 1,
                PIX - 1,
                c,
            )

    # Palette
    for i in range(24):
        px = PALETTE_X + (i % PALETTE_COLS) * (SWATCH_W + SWATCH_GAP_X)
        py = PALETTE_Y + (i // PALETTE_COLS) * (SWATCH_H + SWATCH_GAP_Y)
        draw_logical_color_rect(px, py, SWATCH_W, SWATCH_H, i)
        if i == selected_color and not erase_mode:
            pyxel.rectb(px - 1, py - 1, SWATCH_W + 2, SWATCH_H + 2, 7)

    # Tool buttons
    pyxel.rect(PALETTE_X, TOOL_Y, 52, 18, 13 if erase_mode else 5)
    pyxel.text(PALETTE_X + 6, TOOL_Y + 6, "ERASE", 0)

    pyxel.rect(PALETTE_X, TOOL_Y + 24, 52, 18, 10)
    pyxel.text(PALETTE_X + 5, TOOL_Y + 30, "EXPORT", 0)

    lamp_on = export_lamp_frames > 0 and (pyxel.frame_count // 2) % 2 == 0
    lamp_color = 11 if lamp_on else 2
    pyxel.circb(PALETTE_X + 60, TOOL_Y + 33, 4, 7)
    pyxel.circ(PALETTE_X + 60, TOOL_Y + 33, 3, lamp_color)

    pyxel.rect(PALETTE_X, TOOL_Y + 48, 52, 18, 8)
    pyxel.text(PALETTE_X + 9, TOOL_Y + 54, "CLEAR", 0)

    if ai_phase == "judging":
        c = 10 if (ai_anim_frame // 4) % 2 == 0 else 7
        pyxel.circb(PALETTE_X + 60, TOOL_Y + 57, 5, c)
        pyxel.text(PALETTE_X - 4, TOOL_Y + 62, "AI JUDGING...", c)
    elif ai_phase == "success":
        c = 6 if (ai_anim_frame // 3) % 2 == 0 else 11
        pyxel.circ(PALETTE_X + 60, TOOL_Y + 57, 4, c)
        pyxel.text(PALETTE_X - 8, TOOL_Y + 62, "CORRECT!", c)
    elif ai_phase == "fail":
        c = 3 if (ai_anim_frame // 3) % 2 == 0 else 8
        pyxel.circb(PALETTE_X + 60, TOOL_Y + 57, 5, c)
        pyxel.line(PALETTE_X + 56, TOOL_Y + 53, PALETTE_X + 64, TOOL_Y + 61, c)
        pyxel.line(PALETTE_X + 64, TOOL_Y + 53, PALETTE_X + 56, TOOL_Y + 61, c)
        pyxel.text(PALETTE_X - 10, TOOL_Y + 62, "NOT YET", c)

    if export_lamp_frames > 0:
        pyxel.text(PALETTE_X, TOOL_Y + 68, "SAVING...", 11)
    elif ai_phase == "judging":
        pyxel.text(PALETTE_X, TOOL_Y + 68, "JUDGING...", 10)
    elif ai_phase == "success":
        pyxel.text(PALETTE_X, TOOL_Y + 68, "CORRECT!", 6)
    elif ai_phase == "fail":
        pyxel.text(PALETTE_X, TOOL_Y + 68, "TRY AGAIN", 3)
    else:
        pyxel.text(PALETTE_X, TOOL_Y + 68, "READY", 6)

    used = 0
    for p in pixels:
        if p != 0:
            used += 1
    pyxel.text(CANVAS_X, CANVAS_Y + H * PIX + 6, f"USED DOTS: {used}", 7)
    pyxel.text(CANVAS_X, CANVAS_Y + H * PIX + 12, "TOUCH/DRAG TO DRAW", 7)


pyxel.init(310, 220, title="GAHAKU EXAM")
pyxel.mouse(True)
setup_palette()
export_payload()
pyxel.run(update, draw)
"""

WORD_BANK = [
    "apple",
    "banana",
    "car",
    "cat",
    "dog",
    "fish",
    "house",
    "tree",
    "flower",
    "bird",
    "chair",
    "cup",
    "book",
    "clock",
    "star",
    "moon",
    "sun",
    "heart",
    "shoe",
    "hat",
    "pizza",
    "cake",
    "train",
    "boat",
    "plane",
    "phone",
    "key",
    "camera",
    "bottle",
    "leaf",
    "guitar",
    "ball",
    "rabbit",
    "snake",
    "duck",
    "bus",
    "truck",
    "robot",
    "rocket",
    "bridge",
]

VLM_MODEL_CANDIDATES = [
    "Qwen/Qwen3-VL-8B-Instruct",
    "Qwen/Qwen2.5-VL-3B-Instruct",
    "Qwen/Qwen2.5-VL-7B-Instruct",
    "Qwen/Qwen2-VL-2B-Instruct",
    "Qwen/Qwen2-VL-7B-Instruct",
    "meta-llama/Llama-3.2-11B-Vision-Instruct",
]
ZERO_SHOT_MODEL_CANDIDATES = [
    "openai/clip-vit-large-patch14",
    "openai/clip-vit-base-patch32",
]
IMAGE_CLASSIFICATION_FALLBACK_MODEL = "google/vit-base-patch16-224"
BASE_SCORE = 1000
DOT_PENALTY = 8
MISS_PENALTY = 120
ENABLE_OAUTH_UI = bool(os.getenv("SPACE_ID"))
MODEL_TASK_SUPPORT_CACHE: dict[tuple[str, str], bool] = {}

I18N_TEXTS = {
    "ja": {
        "intro_md": (
            "# GAHAKU EXAM\n\n"
            "## 与えられたお題に従ってドット絵を描き、AIに何の絵か当てさせるゲーム\n"
            "1. 画面下、ラウンド情報の”お題”に従ってPyxelキャンバスでドット絵を描く  \n"
            "2. Pyxel内の **EXPORT** を押す  \n"
            "3. 下の **Python側へ取り込む** を押して受信を確認 \n"
            "4. **AIで判定する** を押し、AIの判定Top3にお題が入っていれば成功！ \n"
            "5. 使ったドットが少ないほど高得点！！"
            
        ),
        "pyxel_help": "Pyxelキャンバス右側の <b>EXPORT</b> を押すと、描画データが window.__PYXEL_PAYLOAD に保存されます。",
        "language_label": "言語 / Language",
        "language_ja": "日本語",
        "language_en": "English",
        "payload_label": "payload(JSON)",
        "sync_btn": "Python側へ取り込む",
        "judge_btn": "AIで判定する",
        "next_btn": "次のラウンド",
        "status_label": "受信結果",
        "status_ready": "準備完了。Pyxelで描いてEXPORTしてください。",
        "raw_json_label": "受信JSON",
        "top5_label": "AI上位5",
        "round_info_title": "## ラウンド情報",
        "round_result_ongoing": "挑戦中",
        "round_result_success": "成功",
        "round_result_fail": "失敗",
        "round_field_round": "Round",
        "round_field_target": "お題",
        "round_field_state": "状態",
        "round_field_miss": "ミス",
        "round_field_used_dots": "使用ドット",
        "round_field_round_score": "ラウンドスコア",
        "round_field_total_score": "合計スコア",
        "err_no_data": "データがまだありません。Pyxel側で EXPORT を押してから再取得してください。",
        "err_parse_json": "JSONの解析に失敗しました",
        "err_invalid_data": "受信データ形式が不正です。",
        "msg_synced": "Python側で受信しました。",
        "msg_new_round": "新しいラウンドを開始しました。Pyxelで絵を描いてEXPORTしてください。",
        "msg_round_finished": "このラウンドは終了済みです。次のラウンドを開始してください。",
        "err_image_convert": "画像変換に失敗しました",
        "err_auth": "AI判定には認証が必要です。Space上では左サイドバーのLogin、ローカルではHF_TOKEN環境変数を設定してください。",
        "err_api": "AI判定APIエラー",
        "err_ai_parse": "AI出力の解析に失敗しました。もう一度判定してください。",
        "msg_correct": "正解です。",
        "msg_incorrect": "不正解です。",
        "msg_game_over": "不正解（3回目）。ゲームオーバーです。",
        "msg_press_next": "次のラウンドを押して再挑戦してください。",
    },
    "en": {
        "intro_md": (
            "# GAHAKU EXAM\n\n"
            "## A game where you draw pixel art from a given target and let AI guess what it is\n"
            "1. Follow the target shown in Round Info and draw pixel art on the Pyxel canvas  \n"
            "2. Press **EXPORT** inside Pyxel  \n"
            "3. Click **Sync from Pyxel** below to confirm receipt  \n"
            "4. Click **Judge with AI**. If the target appears in AI Top 3, you clear the round!  \n"
            "5. Fewer used dots means a higher score!!"
        ),
        "pyxel_help": "Press <b>EXPORT</b> on the right side of the Pyxel canvas to store drawing data in window.__PYXEL_PAYLOAD.",
        "language_label": "言語 / Language",
        "language_ja": "日本語",
        "language_en": "English",
        "payload_label": "payload(JSON)",
        "sync_btn": "Sync from Pyxel",
        "judge_btn": "Judge with AI",
        "next_btn": "Next Round",
        "status_label": "Status",
        "status_ready": "Ready. Draw in Pyxel and press EXPORT.",
        "raw_json_label": "Received JSON",
        "top5_label": "AI Top 5",
        "round_info_title": "## Round Info",
        "round_result_ongoing": "In Progress",
        "round_result_success": "Success",
        "round_result_fail": "Fail",
        "round_field_round": "Round",
        "round_field_target": "Target",
        "round_field_state": "State",
        "round_field_miss": "Miss",
        "round_field_used_dots": "Used dots",
        "round_field_round_score": "Round score",
        "round_field_total_score": "Total score",
        "err_no_data": "No data yet. Press EXPORT in Pyxel and try again.",
        "err_parse_json": "Failed to parse JSON",
        "err_invalid_data": "Invalid payload format.",
        "msg_synced": "Received on Python side.",
        "msg_new_round": "Started a new round. Draw in Pyxel and press EXPORT.",
        "msg_round_finished": "This round is already finished. Start the next round.",
        "err_image_convert": "Failed to convert image",
        "err_auth": "Authentication is required for AI judging. In Spaces, use Login in the left sidebar; locally, set the HF_TOKEN environment variable.",
        "err_api": "AI judge API error",
        "err_ai_parse": "Failed to parse AI output. Please try judging again.",
        "msg_correct": "Correct!",
        "msg_incorrect": "Incorrect.",
        "msg_game_over": "Incorrect (3rd miss). Game over.",
        "msg_press_next": "Press Next Round to try again.",
    },
}

logger = logging.getLogger("gahaku_exam")
if not logger.handlers:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")


@spaces.GPU
def zerogpu_startup_probe():
    # Satisfies ZeroGPU startup validation without affecting current API-based flow.
    return None


def normalize_label(text: str) -> str:
    text = text.strip().lower()
    text = re.sub(r"[^a-z0-9\s]", " ", text)
    text = re.sub(r"\s+", " ", text)
    return text


def normalize_lang(lang: str | None) -> str:
    return "en" if str(lang or "").lower() == "en" else "ja"


def tr(lang: str | None, key: str) -> str:
    lang_n = normalize_lang(lang)
    return I18N_TEXTS[lang_n][key]


def label_matches_target(label: str, target_word: str) -> bool:
    label_n = normalize_label(label)
    target_n = normalize_label(target_word)
    if not target_n:
        return False
    if label_n == target_n:
        return True
    return target_n in label_n.split(" ")


def parse_payload(payload_text: str, lang: str = "ja"):
    if not payload_text:
        return None, tr(lang, "err_no_data")

    try:
        data = json.loads(payload_text)
    except json.JSONDecodeError as exc:
        return None, f"{tr(lang, 'err_parse_json')}: {exc}"

    pixels = data.get("pixels", [])
    width = int(data.get("width", 0))
    height = int(data.get("height", 0))
    palette_hex = data.get("palette_hex", [])

    if not isinstance(pixels, list) or width <= 0 or height <= 0:
        return None, tr(lang, "err_invalid_data")
    if len(pixels) != width * height:
        return None, tr(lang, "err_invalid_data")
    if not isinstance(palette_hex, list) or len(palette_hex) < 24:
        return None, tr(lang, "err_invalid_data")

    return data, None


def payload_to_png_bytes(payload_data: dict) -> bytes:
    from PIL import Image

    width = int(payload_data["width"])
    height = int(payload_data["height"])
    pixels = payload_data["pixels"]
    palette_hex = payload_data["palette_hex"]

    image = Image.new("RGB", (width, height), (255, 255, 255))
    pix = image.load()

    for y in range(height):
        for x in range(width):
            idx = int(pixels[y * width + x])
            idx = max(0, min(idx, len(palette_hex) - 1))
            c = int(palette_hex[idx])
            r = (c >> 16) & 0xFF
            g = (c >> 8) & 0xFF
            b = c & 0xFF
            pix[x, y] = (r, g, b)

    upscaled = image.resize((224, 224), Image.Resampling.NEAREST)
    buf = io.BytesIO()
    upscaled.save(buf, format="PNG")
    return buf.getvalue()


def score_from_state(state: dict) -> int:
    used_dots = int(state.get("used_dots", 0))
    miss_count = int(state.get("miss_count", 0))
    return max(0, BASE_SCORE - DOT_PENALTY * used_dots - MISS_PENALTY * miss_count)


def extract_text_from_chat_response(response) -> str:
    if not hasattr(response, "choices") or not response.choices:
        return ""

    message = getattr(response.choices[0], "message", None)
    if message is None:
        return ""

    content = getattr(message, "content", "")
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        chunks = []
        for item in content:
            if isinstance(item, dict):
                txt = item.get("text")
                if txt:
                    chunks.append(str(txt))
            else:
                txt = getattr(item, "text", None)
                if txt:
                    chunks.append(str(txt))
        return "\n".join(chunks)

    return str(content)


def parse_vlm_guesses(raw_text: str) -> list[str]:
    raw_text = raw_text.strip()
    if not raw_text:
        return []

    # Prefer strict JSON extraction when wrapped with commentary.
    json_match = re.search(r"\{[\s\S]*\}", raw_text)
    candidates = [raw_text]
    if json_match:
        candidates.insert(0, json_match.group(0))

    for candidate in candidates:
        try:
            obj = json.loads(candidate)
        except Exception:
            continue
        guesses = obj.get("guesses") if isinstance(obj, dict) else None
        if isinstance(guesses, list):
            cleaned = []
            for g in guesses:
                label = normalize_label(str(g))
                if label:
                    cleaned.append(label)
            if cleaned:
                return cleaned[:5]

    # Fallback: line/comma separated text
    fallback = []
    for part in re.split(r"[\n,;]+", raw_text):
        label = normalize_label(part)
        if label:
            fallback.append(label)
    return fallback[:5]


def is_model_not_supported_error(detail: str) -> bool:
    text = detail.lower()
    return (
        "model_not_supported" in text
        or "requested model" in text and "not supported" in text
        or "provider" in text and "not support" in text
    )


def has_live_provider_for_task(model_name: str, task: str, token: str | None) -> bool:
    cache_key = (model_name, task)
    if cache_key in MODEL_TASK_SUPPORT_CACHE:
        return MODEL_TASK_SUPPORT_CACHE[cache_key]

    if not token:
        MODEL_TASK_SUPPORT_CACHE[cache_key] = False
        return False

    # Keep '/' unescaped because HF model-info endpoint expects namespace/model path segments.
    encoded_model = urllib.parse.quote(model_name, safe="/")
    url = f"https://huggingface.co/api/models/{encoded_model}?expand=inferenceProviderMapping"
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"})

    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            payload = json.loads(resp.read().decode("utf-8"))
    except Exception as exc:
        logger.warning("provider_mapping_check_failed model=%s task=%s detail=%s", model_name, task, exc)
        MODEL_TASK_SUPPORT_CACHE[cache_key] = False
        return False

    mapping = payload.get("inferenceProviderMapping")
    if not isinstance(mapping, dict):
        MODEL_TASK_SUPPORT_CACHE[cache_key] = False
        return False

    for info in mapping.values():
        if not isinstance(info, dict):
            continue
        if str(info.get("status", "")).lower() == "live" and str(info.get("task", "")).lower() == task:
            MODEL_TASK_SUPPORT_CACHE[cache_key] = True
            return True

    MODEL_TASK_SUPPORT_CACHE[cache_key] = False
    return False


def prediction_item_to_dict(item) -> dict | None:
    label = None
    score = None

    if isinstance(item, dict):
        label = item.get("label")
        score = item.get("score")
    else:
        label = getattr(item, "label", None)
        score = getattr(item, "score", None)

    label_n = normalize_label(str(label or ""))
    if not label_n:
        return None

    try:
        score_f = float(score) if score is not None else 0.0
    except Exception:
        score_f = 0.0

    return {"label": label_n, "score": round(score_f, 4)}


def append_top_predictions(top5: list[dict], predictions, max_items: int = 5):
    for item in predictions or []:
        row = prediction_item_to_dict(item)
        if row is not None:
            top5.append(row)
        if len(top5) >= max_items:
            break


def create_new_round(previous_state: dict | None = None) -> dict:
    previous_state = previous_state or {}
    next_round_number = int(previous_state.get("round_number", 0)) + 1
    total_score = int(previous_state.get("total_score", 0))
    prev_word = previous_state.get("target_word")

    candidates = WORD_BANK
    if prev_word in WORD_BANK and len(WORD_BANK) > 1:
        candidates = [w for w in WORD_BANK if w != prev_word]

    return {
        "round_id": int(time.time() * 1000),
        "round_number": next_round_number,
        "target_word": random.choice(candidates),
        "miss_count": 0,
        "round_result": "ongoing",
        "used_dots": 0,
        "round_score": 0,
        "total_score": total_score,
        "last_top5_predictions": [],
    }


def format_round_info(state: dict, lang: str = "ja") -> str:
    result_map = {
        "ongoing": tr(lang, "round_result_ongoing"),
        "success": tr(lang, "round_result_success"),
        "fail": tr(lang, "round_result_fail"),
    }
    result_label = result_map.get(state.get("round_result"), tr(lang, "round_result_ongoing"))

    return (
        f"{tr(lang, 'round_info_title')}\n"
        f"- {tr(lang, 'round_field_round')}: **{state.get('round_number', 1)}**  \n"
        f"- {tr(lang, 'round_field_target')}: **{state.get('target_word', '-')}**  \n"
        f"- {tr(lang, 'round_field_state')}: **{result_label}**  \n"
        f"- {tr(lang, 'round_field_miss')}: **{state.get('miss_count', 0)} / 3**  \n"
        f"- {tr(lang, 'round_field_used_dots')}: **{state.get('used_dots', 0)}**  \n"
        f"- {tr(lang, 'round_field_round_score')}: **{state.get('round_score', 0)}**  \n"
        f"- {tr(lang, 'round_field_total_score')}: **{state.get('total_score', 0)}**"
    )


def sync_payload_to_round(payload_text: str, round_state: dict, lang: str):
    round_state = round_state or create_new_round()
    lang_n = normalize_lang(lang)

    if not payload_text:
        message = tr(lang_n, "err_no_data")
        return message, round_state, format_round_info(round_state, lang_n)

    try:
        data = json.loads(payload_text)
    except json.JSONDecodeError as exc:
        return f"{tr(lang_n, 'err_parse_json')}: {exc}", round_state, format_round_info(round_state, lang_n)

    pixels = data.get("pixels", [])
    width = int(data.get("width", 0))
    height = int(data.get("height", 0))

    if not isinstance(pixels, list) or width <= 0 or height <= 0:
        return tr(lang_n, "err_invalid_data"), round_state, format_round_info(round_state, lang_n)

    used_dots = sum(1 for p in pixels if p != 0)
    unique_colors = sorted({int(p) for p in pixels if isinstance(p, int) and p != 0})
    round_state["used_dots"] = used_dots

    summary = (
        f"{tr(lang_n, 'msg_synced')}"
        f"\n- target_word: {round_state.get('target_word', '-') }"
        f"\n- size: {width}x{height}"
        f"\n- used_dots: {used_dots}"
        f"\n- unique_colors: {unique_colors}"
    )
    return summary, round_state, format_round_info(round_state, lang_n)


def start_next_round(round_state: dict, lang: str):
    next_state = create_new_round(round_state)
    lang_n = normalize_lang(lang)
    info = format_round_info(next_state, lang_n)
    message = tr(lang_n, "msg_new_round")
    return next_state, info, message


def judge_with_ai_core(payload_text: str, round_state: dict, token: str | None, lang: str):
    round_state = round_state or create_new_round()
    lang_n = normalize_lang(lang)

    if round_state.get("round_result") in {"success", "fail"}:
        message = tr(lang_n, "msg_round_finished")
        return message, round_state.get("last_top5_predictions", []), round_state, format_round_info(round_state, lang_n)

    payload_data, parse_error = parse_payload(payload_text, lang_n)
    if parse_error:
        logger.info("judge_parse_error round=%s error=%s", round_state.get("round_id"), parse_error)
        return parse_error, [], round_state, format_round_info(round_state, lang_n)

    pixels = payload_data["pixels"]
    round_state["used_dots"] = sum(1 for p in pixels if int(p) != 0)
    logger.info(
        "judge_start round=%s word=%s used_dots=%s",
        round_state.get("round_id"),
        round_state.get("target_word"),
        round_state.get("used_dots"),
    )

    try:
        image_bytes = payload_to_png_bytes(payload_data)
    except Exception as exc:
        return f"{tr(lang_n, 'err_image_convert')}: {exc}", [], round_state, format_round_info(round_state, lang_n)

    image_b64 = base64.b64encode(image_bytes).decode("ascii")
    image_url = f"data:image/png;base64,{image_b64}"

    top5 = []
    last_error_detail = ""
    unsupported_vlm_models = []

    for model_name in VLM_MODEL_CANDIDATES:
        if not has_live_provider_for_task(model_name, "conversational", token):
            logger.info("vlm_skip_no_provider round=%s model=%s", round_state.get("round_id"), model_name)
            unsupported_vlm_models.append(model_name)
            continue
        logger.info("vlm_try round=%s model=%s", round_state.get("round_id"), model_name)
        try:
            client = InferenceClient(token=token, model=model_name)
            response = client.chat_completion(
                model=model_name,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You are a visual guesser for tiny pixel art. "
                            "The image is a 16x16 user-drawn icon. "
                            "Do not explain. Return JSON only."
                        ),
                    },
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": (
                                    "Guess what object this 16x16 pixel art depicts. "
                                    "Return exactly 5 likely guesses in English, most likely first. "
                                    "Output JSON only in this format: "
                                    '{"guesses": ["word1", "word2", "word3", "word4", "word5"]}'
                                ),
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": image_url},
                            },
                        ],
                    },
                ],
                max_tokens=220,
                temperature=0.2,
            )
        except Exception as exc:
            detail = str(exc)
            last_error_detail = detail
            if "401" in detail or "Unauthorized" in detail:
                logger.warning("vlm_auth_error round=%s model=%s", round_state.get("round_id"), model_name)
                message = tr(lang_n, "err_auth")
                return message, [], round_state, format_round_info(round_state, lang_n)
            if is_model_not_supported_error(detail):
                logger.warning("vlm_not_supported round=%s model=%s detail=%s", round_state.get("round_id"), model_name, detail)
                unsupported_vlm_models.append(model_name)
                continue
            logger.exception("vlm_error round=%s model=%s", round_state.get("round_id"), model_name)
            return f"{tr(lang_n, 'err_api')}: {exc}", [], round_state, format_round_info(round_state, lang_n)

        raw_text = extract_text_from_chat_response(response)
        guesses = parse_vlm_guesses(raw_text)
        if not guesses:
            last_error_detail = f"VLM応答解析失敗 ({model_name})"
            logger.warning("vlm_parse_failed round=%s model=%s", round_state.get("round_id"), model_name)
            continue

        logger.info("vlm_success round=%s model=%s top1=%s", round_state.get("round_id"), model_name, guesses[0])

        for i, label in enumerate(guesses[:5]):
            # VLMは確率を返さないため、順位ベースの擬似スコアを表示する。
            top5.append({"label": label, "score": round(max(0.0, 1.0 - i * 0.1), 4)})
        break

    if not top5:
        for model_name in ZERO_SHOT_MODEL_CANDIDATES:
            if not has_live_provider_for_task(model_name, "zero-shot-image-classification", token):
                logger.info("zero_shot_skip_no_provider round=%s model=%s", round_state.get("round_id"), model_name)
                continue
            logger.info("zero_shot_try round=%s model=%s", round_state.get("round_id"), model_name)
            try:
                zsc_client = InferenceClient(token=token, model=model_name)
                zsc_predictions = zsc_client.zero_shot_image_classification(
                    image=image_bytes,
                    candidate_labels=WORD_BANK,
                    model=model_name,
                    hypothesis_template="This icon is a {}.",
                )
                append_top_predictions(top5, zsc_predictions, max_items=5)
                if top5:
                    logger.info(
                        "zero_shot_success round=%s model=%s top1=%s",
                        round_state.get("round_id"),
                        model_name,
                        top5[0]["label"],
                    )
                    break
                logger.warning("zero_shot_empty round=%s model=%s", round_state.get("round_id"), model_name)
            except StopIteration:
                logger.warning("zero_shot_not_supported round=%s model=%s detail=no_provider_mapping", round_state.get("round_id"), model_name)
                continue
            except Exception as exc:
                detail = str(exc)
                if is_model_not_supported_error(detail):
                    logger.warning(
                        "zero_shot_not_supported round=%s model=%s detail=%s",
                        round_state.get("round_id"),
                        model_name,
                        detail,
                    )
                    continue
                if "401" in detail or "Unauthorized" in detail:
                    logger.warning("zero_shot_auth_error round=%s model=%s", round_state.get("round_id"), model_name)
                    message = tr(lang_n, "err_auth")
                    return message, [], round_state, format_round_info(round_state, lang_n)
                logger.exception("zero_shot_error round=%s model=%s", round_state.get("round_id"), model_name)

    if not top5:
        logger.info(
            "fallback_to_image_classification round=%s model=%s unsupported_vlm=%s",
            round_state.get("round_id"),
            IMAGE_CLASSIFICATION_FALLBACK_MODEL,
            unsupported_vlm_models,
        )
        try:
            clf_client = InferenceClient(token=token, model=IMAGE_CLASSIFICATION_FALLBACK_MODEL)
            predictions = clf_client.image_classification(
                image=image_bytes,
                model=IMAGE_CLASSIFICATION_FALLBACK_MODEL,
                top_k=5,
            )
            append_top_predictions(top5, predictions, max_items=5)
            if top5:
                logger.info(
                    "image_classification_success round=%s model=%s top1=%s",
                    round_state.get("round_id"),
                    IMAGE_CLASSIFICATION_FALLBACK_MODEL,
                    top5[0]["label"],
                )
        except Exception as exc:
            detail = str(exc)
            if "401" in detail or "Unauthorized" in detail:
                logger.warning("image_classification_auth_error round=%s", round_state.get("round_id"))
                message = tr(lang_n, "err_auth")
                return message, [], round_state, format_round_info(round_state, lang_n)
            logger.exception(
                "image_classification_error round=%s model=%s",
                round_state.get("round_id"),
                IMAGE_CLASSIFICATION_FALLBACK_MODEL,
            )
            fallback_note = ""
            if unsupported_vlm_models:
                fallback_note = f" VLM未対応: {', '.join(unsupported_vlm_models)}"
            return f"{tr(lang_n, 'err_api')}: {exc}.{fallback_note}", [], round_state, format_round_info(round_state, lang_n)

    if not top5:
        logger.warning(
            "judge_no_predictions round=%s unsupported_vlm=%s last_error=%s",
            round_state.get("round_id"),
            unsupported_vlm_models,
            last_error_detail,
        )
        note = ""
        if unsupported_vlm_models:
            note = f" (VLM未対応: {', '.join(unsupported_vlm_models)})"
        if last_error_detail:
            note = f"{note} 詳細: {last_error_detail}"
        return f"{tr(lang_n, 'err_ai_parse')}{note}", [], round_state, format_round_info(round_state, lang_n)

    round_state["last_top5_predictions"] = top5

    target = round_state.get("target_word", "")
    # Difficulty tweak: count as correct only when target is within top 3 guesses.
    is_correct = any(label_matches_target(pred["label"], target) for pred in top5[:3])

    if is_correct:
        round_state["round_result"] = "success"
        round_state["round_score"] = score_from_state(round_state)
        round_state["total_score"] = int(round_state.get("total_score", 0)) + int(round_state["round_score"])
        logger.info(
            "judge_result round=%s result=success misses=%s score=%s",
            round_state.get("round_id"),
            round_state.get("miss_count", 0),
            round_state.get("round_score", 0),
        )
        message = (
            f"{tr(lang_n, 'msg_correct')}\n"
            f"- target_word: {target}\n"
            f"- used_dots: {round_state['used_dots']}\n"
            f"- round_score: {round_state['round_score']}"
        )
    else:
        round_state["miss_count"] = int(round_state.get("miss_count", 0)) + 1
        if round_state["miss_count"] >= 3:
            round_state["round_result"] = "fail"
            logger.info(
                "judge_result round=%s result=fail misses=%s",
                round_state.get("round_id"),
                round_state.get("miss_count", 0),
            )
            message = (
                f"{tr(lang_n, 'msg_game_over')}\n"
                f"- target_word: {target}\n"
                f"{tr(lang_n, 'msg_press_next')}"
            )
        else:
            logger.info(
                "judge_result round=%s result=miss misses=%s",
                round_state.get("round_id"),
                round_state.get("miss_count", 0),
            )
            message = (
                f"{tr(lang_n, 'msg_incorrect')}\n"
                f"- target_word: {target}\n"
                f"- miss_count: {round_state['miss_count']} / 3"
            )

    return message, top5, round_state, format_round_info(round_state, lang_n)


def judge_with_ai_local(payload_text: str, round_state: dict, lang: str):
    return judge_with_ai_core(payload_text, round_state, os.getenv("HF_TOKEN"), lang)


def judge_with_ai_space(
    payload_text: str,
    round_state: dict,
    lang: str,
    hf_token: gr.OAuthToken | None = None,
):
    token = hf_token.token if hf_token and getattr(hf_token, "token", None) else os.getenv("HF_TOKEN")
    return judge_with_ai_core(payload_text, round_state, token, lang)


def build_pyxel_html() -> str:
    # Ensure JS-side animation state starts clean on each page load.
    init_js = '<script>window.__AI_JUDGE_STATE = "idle";</script>'
    escaped = (
        PYXEL_SCRIPT.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )
    return f"""
{init_js}
<div style="border:1px solid #d1d5db;border-radius:12px;padding:8px;background:#111827;">
    <div id="pyxel-stage" style="position:relative;width:min(100%,620px);aspect-ratio:31/22;margin:0 auto;overflow:hidden;">
        <div id="pyxel-screen" style="width:100%;height:100%;"></div>
    </div>
  <pyxel-run script="{escaped}"></pyxel-run>
</div>
"""


def build_pyxel_help_markdown(lang: str) -> str:
    return f"<p style=\"margin-top:8px;font-size:0.9rem;color:#374151;\">{tr(lang, 'pyxel_help')}</p>"


def apply_language(lang: str, round_state: dict):
    lang_n = normalize_lang(lang)
    state = round_state or create_new_round()
    return (
        lang_n,
        gr.update(value=tr(lang_n, "intro_md")),
        gr.update(value=format_round_info(state, lang_n)),
        gr.update(label=tr(lang_n, "payload_label")),
        gr.update(value=tr(lang_n, "sync_btn")),
        gr.update(value=tr(lang_n, "judge_btn")),
        gr.update(value=tr(lang_n, "next_btn")),
        gr.update(label=tr(lang_n, "status_label")),
        gr.update(label=tr(lang_n, "top5_label")),
    )


head_html = """
<link rel="apple-touch-icon" href="/favicon.ico">
<link rel="apple-touch-icon-precomposed" href="/favicon.ico">
<script src="https://cdn.jsdelivr.net/gh/kitao/pyxel@v2.9.5/wasm/pyxel.js"></script>
"""


with gr.Blocks(title="GAHAKU EXAM") as demo:
    if ENABLE_OAUTH_UI:
        with gr.Sidebar():
            gr.LoginButton()

    default_lang = "en"
    lang_state = gr.State(default_lang)

    with gr.Row(equal_height=True):
        with gr.Column(scale=8):
            intro_md = gr.Markdown(tr(default_lang, "intro_md"))
        with gr.Column(scale=2, min_width=220):
            lang_toggle = gr.Radio(
                choices=[
                    (tr(default_lang, "language_ja"), "ja"),
                    (tr(default_lang, "language_en"), "en"),
                ],
                value=default_lang,
                label=tr(default_lang, "language_label"),
            )

    pyxel_html = gr.HTML(build_pyxel_html())
    initial_round_state = create_new_round()
    round_state = gr.State(initial_round_state)
    round_info = gr.Markdown(format_round_info(initial_round_state, default_lang))

    payload_box = gr.Textbox(label=tr(default_lang, "payload_label"), elem_id="pyxel-payload", visible=False)
    with gr.Row():
        sync_btn = gr.Button(tr(default_lang, "sync_btn"), variant="primary")
        judge_btn = gr.Button(tr(default_lang, "judge_btn"), variant="primary")
        next_round_btn = gr.Button(tr(default_lang, "next_btn"), variant="secondary")

    status_box = gr.Textbox(label=tr(default_lang, "status_label"), value=tr(default_lang, "status_ready"))
    top5_box = gr.JSON(label=tr(default_lang, "top5_label"))

    lang_toggle.change(
        fn=apply_language,
        inputs=[lang_toggle, round_state],
        outputs=[
            lang_state,
            intro_md,
            round_info,
            payload_box,
            sync_btn,
            judge_btn,
            next_round_btn,
            status_box,
            top5_box,
        ],
    )

    sync_btn.click(
        fn=sync_payload_to_round,
        inputs=[payload_box, round_state, lang_state],
                outputs=[status_box, round_state, round_info],
        js="""
        (current, state) => {
          const payload = window.__PYXEL_PAYLOAD;
          if (typeof payload === "string" && payload.length > 0) {
            return [payload, state];
          }
          return [current, state];
        }
        """,
    )

    next_round_btn.click(
        fn=start_next_round,
        inputs=[round_state, lang_state],
        outputs=[round_state, round_info, status_box],
                js="""
                (state, lang) => {
                    window.__AI_JUDGE_STATE = "idle";
                    return [state, lang];
                }
                """,
    )

    if ENABLE_OAUTH_UI:
        judge_btn.click(
            fn=judge_with_ai_space,
            inputs=[payload_box, round_state, lang_state],
            outputs=[status_box, top5_box, round_state, round_info],
            js="""
            (current, state) => {
                            window.__AI_JUDGE_STATE = "judging";
              const payload = window.__PYXEL_PAYLOAD;
              if (typeof payload === "string" && payload.length > 0) {
                return [payload, state];
              }
              return [current, state];
            }
            """,
        )
    else:
        judge_btn.click(
            fn=judge_with_ai_local,
            inputs=[payload_box, round_state, lang_state],
            outputs=[status_box, top5_box, round_state, round_info],
            js="""
            (current, state) => {
                            window.__AI_JUDGE_STATE = "judging";
              const payload = window.__PYXEL_PAYLOAD;
              if (typeof payload === "string" && payload.length > 0) {
                return [payload, state];
              }
              return [current, state];
            }
            """,
        )

        status_box.change(
                fn=None,
                inputs=[status_box],
                outputs=[],
                js="""
                (status) => {
                    const text = String(status || "");
                    if (text.includes("正解です") || text.includes("Correct")) {
                        window.__AI_JUDGE_STATE = "success";
                    } else if (
                        text.includes("不正解") ||
                        text.includes("ゲームオーバー") ||
                        text.includes("Incorrect") ||
                        text.includes("Game over")
                    ) {
                        window.__AI_JUDGE_STATE = "fail";
                    }
                    return [];
                }
                """,
        )


if __name__ == "__main__":
    demo.launch(head=head_html, ssr_mode=False)