"""Doodle Duel — human draws, Qwen (vision) guesses in real time.

Run:   uv run doodle_duel/app.py            (mock mode)
       MODEL_BASE_URL=... uv run ...         (needs a vision model endpoint)
"""
from __future__ import annotations

import os
import threading
import time

import gradio as gr
import numpy as np

import config
import game
import visuals
import vision_client
from game import DoodleState


# ── Canvas helpers ────────────────────────────────────────────────────────────

def _blank_canvas():
    white = np.full((config.CANVAS_H, config.CANVAS_W, 3), 255, np.uint8)
    return {"background": white, "layers": [], "composite": white}


# ── render() returns ALL UI outputs ──────────────────────────────────────────
#
# Output order (13 total):
#   [0] topbar  [1] ai_bubble  [2] thinking  [3] result
#   [4..6] opt[0..2]  (word-choice buttons)
#   [7] easy_btn  [8] medium_btn  [9] hard_btn  (mode buttons)
#   [10] clear_btn  [11] hint_btn
#   [12] canvas
#
N_RENDER = 13

_DIFF_LABELS = {
    "easy":   "🟢 Easy",
    "medium": "🟡 Medium",
    "hard":   "🔴 Hard",
}

def render(state: DoodleState, clear_canvas: bool = False):
    p = state.status
    choosing = p == "choosing"
    drawing  = p == "drawing"
    over     = p in ("won", "lost")
    idle     = p == "idle"
    show_mode = idle or over        # difficulty buttons visible between rounds

    opts = []
    for i in range(config.WORD_CHOICES):
        if choosing and i < len(state.options):
            opts.append(gr.update(value=state.options[i], visible=True))
        else:
            opts.append(gr.update(visible=False))

    def _mode_btn(diff):
        label = _DIFF_LABELS[diff]
        if over and state.difficulty == diff:
            label = f"{_DIFF_LABELS[diff]}  ↺"
        return gr.update(value=label,
                         variant="primary" if state.difficulty == diff else "secondary",
                         visible=show_mode)

    easy_upd   = _mode_btn("easy")
    medium_upd = _mode_btn("medium")
    hard_upd   = _mode_btn("hard")

    clear_upd = gr.update(visible=drawing)
    hints_left = config.MAX_HINTS - state.hints_used
    hint_upd  = gr.update(value=f"💡  Hint ({hints_left} left)",
                           visible=drawing and hints_left > 0)

    canvas_upd = gr.update(value=_blank_canvas()) if clear_canvas else gr.skip()

    return (
        visuals.render_topbar(state),
        visuals.render_ai(state),
        visuals.render_thinking(state),
        visuals.render_result(state),
        *opts,
        easy_upd, medium_upd, hard_upd,
        clear_upd, hint_upd,
        canvas_upd,
    )


# ── Event handlers ────────────────────────────────────────────────────────────

def _warmup_endpoint():
    threading.Thread(target=vision_client.health, daemon=True).start()


# ── Sequential guess gate ─────────────────────────────────────────────────────
# Only the timer triggers guesses now (drawing just stashes the canvas). This
# module-level, per-session gate makes those guesses STRICTLY SEQUENTIAL: a new
# model request starts only when (a) none is already in flight, and (b) at least
# GUESS_MIN_INTERVAL has passed since the previous one finished. We keep this in
# a module global rather than gr.State because gr.State hands each handler a
# *copy* — concurrent timer ticks wouldn't share an in-flight flag through it,
# which is what let parallel requests through and tripped HF's 429.
_GUESS = {}                       # sid -> {"in_flight": bool, "last": float}
_GUESS_LOCK = threading.Lock()


def _sid(request) -> str:
    return getattr(request, "session_hash", None) or "default"


def _claim_guess(sid, force=False) -> bool:
    """Try to claim the single guess slot for this session. Returns True only if
    no guess is in flight and (unless force) the interval has elapsed; the caller
    MUST call _release_guess(sid) when the request completes."""
    now = time.monotonic()
    with _GUESS_LOCK:
        g = _GUESS.setdefault(sid, {"in_flight": False, "last": 0.0})
        if g["in_flight"]:
            return False
        if not force and now - g["last"] < config.GUESS_MIN_INTERVAL:
            return False
        g["in_flight"] = True
        return True


def _release_guess(sid):
    """Release the slot and restart the interval clock once a response returns."""
    with _GUESS_LOCK:
        g = _GUESS.setdefault(sid, {"in_flight": False, "last": 0.0})
        g["in_flight"] = False
        g["last"] = time.monotonic()


# ── Model status pill ─────────────────────────────────────────────────────────
# The model runs on an HF Inference Endpoint that scales to zero when idle, so
# the first request after a quiet spell triggers a cold start (~1-2 min). The
# pill below tells the player whether the AI is ready, waking, or offline so an
# empty "Analyzing..." panel during a cold start isn't mistaken for a crash.
_health = {"checked": False, "ok": False, "fails": 0}


def _status_pill(text, color, pulse=False):
    dot = (f'<span style="display:inline-block;width:9px;height:9px;border-radius:50%;'
           f'background:{color};margin-right:7px;'
           f'{"animation:dd-pulse 1.1s ease-in-out infinite;" if pulse else ""}"></span>')
    # Fixed-height, single-line wrapper. The pill text flips from the long
    # "Waking the AI…" string to a short "AI ready" exactly once — right when the
    # first prediction lands and the model warms up. If the long text is allowed
    # to wrap, that flip shrinks the pill from two lines to one, shifting the
    # canvas below it UP mid-stroke and leaving a stray diagonal. Reserving a
    # constant height and forbidding wrap keeps the layout perfectly still.
    return (
        '<div style="display:flex;justify-content:center;align-items:center;'
        'height:32px;margin:-4px 0 6px">'
        '<div style="display:inline-flex;align-items:center;white-space:nowrap;'
        'font-size:13px;font-weight:600;'
        'padding:5px 14px;border-radius:999px;background:#ffffff14;'
        f'border:1px solid {color}55;color:#e8e8f0">{dot}{text}</div></div>'
        '<style>@keyframes dd-pulse{0%,100%{opacity:1}50%{opacity:.35}}</style>'
    )


def render_status():
    if config.MOCK_MODE:
        return _status_pill("Mock mode — no model connected", "#a78bfa")
    if not _health["checked"]:
        return _status_pill("Connecting to the AI…", "#f4c04e", pulse=True)
    if _health["ok"]:
        return _status_pill("AI ready", "#46d39a")
    if _health["fails"] <= 6:
        return _status_pill("Waking the AI… (cold start, up to ~2 min)", "#f4c04e", pulse=True)
    return _status_pill("AI offline — retrying…", "#ef5d6c", pulse=True)


def on_health():
    ok, _ = vision_client.health()
    _health["checked"] = True
    _health["ok"] = ok
    _health["fails"] = 0 if ok else _health["fails"] + 1
    return render_status()


def on_load():
    _warmup_endpoint()
    s = DoodleState()
    return (s, *render(s, clear_canvas=True))


def on_mode(s, difficulty):
    """User clicked Easy / Medium / Hard — offer 3 words from that pool."""
    if s is None:
        s = DoodleState()
    game.offer_words(s, difficulty)
    return (s, *render(s, clear_canvas=True))


def on_choose(s, i):
    if s is None:
        return (s, *render(DoodleState(), clear_canvas=True))
    game.choose_word(s, i)
    return (s, *render(s, clear_canvas=True))


def on_capture(s, canvas):
    """Stroke-end (canvas.change): ONLY stash the canvas and update the hidden
    state. This event outputs to `st` ALONE (see binding) — never the canvas or
    any visible component — so a stroke triggers ZERO re-render: no flicker, no
    'reload' feel, and crucially no value pushed back to the Sketchpad (pushing
    a value mid/post-stroke is what left the stray diagonal line and disturbed
    the canvas). The timer reads this stashed canvas to do all the guessing."""
    if s is not None:
        game.stash_canvas(s, canvas)
    return s


def on_timer(s, request: gr.Request):
    """Timer tick: always update the countdown; guess only when we can claim the
    sequential slot (no request in flight + interval elapsed). The guess runs
    synchronously here, but other ticks run concurrently and just render the
    countdown, so a slow / cold-start call never freezes the clock. The player
    keeps getting fresh guesses even while paused — without overlapping calls.
    Outputs everything EXCEPT the canvas, so the timer never disturbs drawing."""
    if s is None:
        return (s, *([gr.skip()] * (N_RENDER - 1)))
    game.tick_time(s)                       # countdown / timeout, no network
    if s.status == "drawing":
        sid = _sid(request)
        if _claim_guess(sid):
            try:
                game.poll_guess(s, s.last_canvas)
            finally:
                _release_guess(sid)
    *other_outs, _ = render(s)              # drop the trailing canvas update
    return (s, *other_outs)


def on_hint(s, canvas, request: gr.Request):
    if s is None:
        return (s, *([gr.skip()] * (N_RENDER - 1)))
    game.use_hint(s)
    game.stash_canvas(s, canvas)            # capture the current drawing
    sid = _sid(request)
    if _claim_guess(sid, force=True):       # guess now (bypass interval, respect in-flight)
        try:
            game.poll_guess(s, canvas)
        finally:
            _release_guess(sid)
    *other_outs, _ = render(s)              # drop the trailing canvas update
    return (s, *other_outs)


# ── Injected JS / CSS ─────────────────────────────────────────────────────────
_css_path = "style.css" if os.path.exists("style.css") else "doodle_duel/style.css"
try:
    with open(_css_path) as _f:
        _CSS_TEXT = _f.read()
except FileNotFoundError:
    _CSS_TEXT = ""

_JS = """
<canvas id="dd-particles"></canvas>
<script src="https://cdn.jsdelivr.net/npm/gsap@3.12.5/dist/gsap.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/canvas-confetti@1.9.3/dist/confetti.browser.min.js"></script>
<script>
// While the pen is down on the drawing canvas, the main thread must stay free:
// any heavy work during a stroke (particle frames, GSAP chip animations, the
// right-panel HTML swap when a guess lands) can drop a pointermove, and Gradio's
// canvas then connects the gap with a straight line — the stray diagonal. We
// detect drawing via pointer events on the sketchpad's canvas and pause the
// decorative animations until the stroke ends.
window.ddDrawing = false;
document.addEventListener('pointerdown', function(e){
  var t = e.target;
  if (t && t.tagName === 'CANVAS' && t.id !== 'dd-particles') window.ddDrawing = true;
}, true);
function ddStopDraw(){ window.ddDrawing = false; }
document.addEventListener('pointerup', ddStopDraw, true);
document.addEventListener('pointercancel', ddStopDraw, true);
window.addEventListener('blur', ddStopDraw);

(function() {
  var c = document.getElementById('dd-particles');
  if (!c) return;
  var ctx = c.getContext('2d');
  var W = 0, H = 0;
  function resize() { W = c.width = window.innerWidth; H = c.height = window.innerHeight; }
  resize();
  window.addEventListener('resize', resize);
  var pts = [];
  for (var i = 0; i < 75; i++) {
    pts.push({ x:Math.random()*W, y:Math.random()*H,
               r:Math.random()*1.8+0.4,
               dx:(Math.random()-0.5)*0.3, dy:(Math.random()-0.5)*0.3,
               hue:Math.random()*60+220, a:Math.random()*0.45+0.1 });
  }
  function draw() {
    requestAnimationFrame(draw);
    if (window.ddDrawing) return;        // freeze particles mid-stroke (keep main thread free)
    ctx.clearRect(0,0,W,H);
    pts.forEach(function(p){
      ctx.beginPath(); ctx.arc(p.x,p.y,p.r,0,Math.PI*2);
      ctx.fillStyle='hsla('+p.hue+',80%,70%,'+p.a+')'; ctx.fill();
      p.x+=p.dx; p.y+=p.dy;
      if(p.x<0||p.x>W) p.dx*=-1;
      if(p.y<0||p.y>H) p.dy*=-1;
    });
  }
  draw();
})();
var _lastWin=false;
function checkWin(){
  var el=document.getElementById('win-result');
  if(el && !_lastWin){
    _lastWin=true;
    if(typeof confetti!=='undefined'){
      confetti({particleCount:150,spread:85,origin:{y:0.6},
                colors:['#a78bfa','#f472b6','#22d3ee','#fbbf24','#34d399']});
    }
  } else if(!el){ _lastWin=false; }
  setTimeout(checkWin,700);
}
checkWin();
function animateNew(){
  if(typeof gsap==='undefined') return;
  if(window.ddDrawing) return;          // don't animate chips/results mid-stroke
  document.querySelectorAll('.result:not(.gd)').forEach(function(el){
    el.classList.add('gd');
    gsap.from(el,{duration:0.5,y:30,opacity:0,ease:'back.out(1.7)'});
  });
  document.querySelectorAll('.chip:not(.gd)').forEach(function(el){
    el.classList.add('gd');
    gsap.from(el,{duration:0.3,scale:0.4,opacity:0,ease:'back.out(2)'});
  });
}
setInterval(animateNew,350);
</script>
"""

_HEAD_HTML = "<style>" + _CSS_TEXT + "</style>" + _JS


# ── Build the Gradio UI ───────────────────────────────────────────────────────

with gr.Blocks(title="Doodle Duel") as demo:
    st = gr.State()

    gr.HTML(_HEAD_HTML)
    gr.HTML('<div id="title">🎨 Doodle Duel</div>'
            '<div id="blurb">Draw a word — the robot tries to guess it in real time!</div>')

    model_status = gr.HTML(render_status())

    topbar = gr.HTML()

    with gr.Row(elem_id="word-choices"):
        opt = [gr.Button("…", visible=False) for _ in range(config.WORD_CHOICES)]

    # Side by side again, but every dimension is frozen (see CSS): both columns
    # are FIXED pixel widths (immune to the scrollbar changing the container
    # width), heights are decoupled, the thoughts panel is a fixed-size box that
    # scrolls internally, and the chip bubble lives BELOW the stage so nothing in
    # the canvas's own column can ever grow. Result: the canvas can't move or
    # rescale when a guess lands.
    with gr.Row(elem_id="stage"):
        with gr.Column(elem_id="canvas-col"):
            canvas = gr.Sketchpad(
                height=config.CANVAS_H, type="numpy", label="", show_label=False,
                canvas_size=(config.CANVAS_W, config.CANVAS_H),
                brush=gr.Brush(default_size=config.BRUSH_SIZE, default_color="#111111"),
            )
        with gr.Column(elem_id="thoughts-col"):
            thinking = gr.HTML()

    ai_bubble = gr.HTML()

    result = gr.HTML()

    with gr.Row(elem_id="controls"):
        easy_btn   = gr.Button("🟢 Easy",   variant="secondary", scale=2)
        medium_btn = gr.Button("🟡 Medium", variant="primary",   scale=2)
        hard_btn   = gr.Button("🔴 Hard",   variant="secondary", scale=2)
        hint_btn   = gr.Button("💡  Hint",  scale=2, visible=False)
        clear_btn  = gr.Button("🧽  Clear", scale=1, visible=False)

    timer = gr.Timer(1.0)          # 1 Hz: countdown shows mm:ss, and halves the
                                   # request rate vs the old 0.5 s tick (eased the 429)
    health_timer = gr.Timer(5.0)   # poll model health (lightweight /health route)

    render_outs = [
        topbar, ai_bubble, thinking, result,
        *opt,
        easy_btn, medium_btn, hard_btn,
        clear_btn, hint_btn,
        canvas,
    ]
    assert len(render_outs) == N_RENDER, f"{len(render_outs)} != {N_RENDER}"

    all_outs = [st, *render_outs]
    # Same outputs minus the canvas (which is render_outs[-1]). Events that must
    # never touch the Sketchpad (timer, hint, stroke-capture) use this so they
    # can't push a value back to the canvas and disturb the drawing.
    all_outs_no_canvas = all_outs[:-1]

    demo.load(on_load, outputs=all_outs)
    demo.load(on_health, outputs=[model_status])
    health_timer.tick(on_health, outputs=[model_status])

    easy_btn.click(lambda s: on_mode(s, "easy"),     inputs=st, outputs=all_outs)
    medium_btn.click(lambda s: on_mode(s, "medium"), inputs=st, outputs=all_outs)
    hard_btn.click(lambda s: on_mode(s, "hard"),     inputs=st, outputs=all_outs)

    clear_btn.click(lambda s: (s, *render(s, clear_canvas=True)),
                    inputs=st, outputs=all_outs)
    hint_btn.click(on_hint, inputs=[st, canvas], outputs=all_outs_no_canvas)

    for i in range(config.WORD_CHOICES):
        opt[i].click(lambda s, i=i: on_choose(s, i), inputs=st, outputs=all_outs)

    # Stroke-end only stashes the canvas into the hidden state — it outputs to
    # `st` ALONE, so a stroke re-renders nothing and never writes back to the
    # canvas. The timer is the sole guesser; it reads the stashed canvas, fires
    # one sequential guess at a time, and outputs everything except the canvas.
    canvas.change(on_capture, inputs=[st, canvas], outputs=[st])
    timer.tick(on_timer, inputs=[st], outputs=all_outs_no_canvas)

    # Allow timer ticks to run concurrently: while one tick is blocked on a slow
    # (or cold-start) model call, the others still render the countdown, so the
    # clock never freezes. The per-session in-flight gate keeps actual model
    # calls sequential regardless of how many ticks run at once.
    demo.queue(default_concurrency_limit=16)


if __name__ == "__main__":
    ok, detail = vision_client.health()
    print(f"🎨 Doodle Duel — model backend: {'OK  ' if ok else 'DOWN '}{detail}")
    port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
    demo.launch(server_name="0.0.0.0", server_port=port,
                css_paths=[_css_path], theme=gr.themes.Base())