"""World Cup Heros — football shot analysis (Gradio, ZeroGPU Space).

One RF-DETR-Seg model detects ball + players + goal; pure geometry/physics derive the
events: was it a goal (+ when), who shot it, their pose & shooting foot, and shot speed.
Upload a clip or pick one of the 12 held-out test examples.

ZeroGPU: only the RF-DETR detection runs inside the @spaces.GPU window; the goal physics,
possession, MediaPipe pose, speed and rendering all run on CPU.

Left:   upload + example clips.
Right:  the clean analysed video (player masks + P# labels + ball trail + goal), then a
        Result table (goal / foot, predicted vs actual), a Shot-speed card, and a
        Pose-Detection card.
"""
import glob
import json
import os

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

from futheros import pipeline as P
from futheros import coach as C
from futheros.render import render, pose_card_img

HF_MODEL = "build-small-hackathon/futheros-rfdetr-seg"   # RF-DETR-Seg-Small checkpoint
RENDER_DIR = "/tmp/app_renders"; os.makedirs(RENDER_DIR, exist_ok=True)
GT = json.load(open("labels/goal_labels.json"))
LEG = json.load(open("labels/leg_labels.json"))
EXAMPLES = sorted(glob.glob("examples/*.mp4"))
EDGE_EMOJIS = ["🏴", "⚽", "🇫🇷", "⚽", "🇧🇷", "⚽", "🇦🇷", "⚽", "🇪🇸", "⚽",
               "🇩🇪", "⚽", "🇵🇹", "⚽", "🇳🇱", "⚽", "🇧🇪", "⚽"]
EDGE_STRIP = "  \n".join(EDGE_EMOJIS * 2)
CONFETTI_JS = """
(g) => {
  if (g === 'goal') {
    const fire = () => {
      window.confetti({particleCount:200, spread:100, origin:{y:0.6}});
      setTimeout(()=>window.confetti({particleCount:150, spread:120, origin:{y:0.5}}), 200);
      setTimeout(()=>window.confetti({particleCount:120, angle:60, spread:80, origin:{x:0}}), 350);
      setTimeout(()=>window.confetti({particleCount:120, angle:120, spread:80, origin:{x:1}}), 350);
    };
    if (!window.confetti) {
      const s = document.createElement('script');
      s.src = 'https://cdn.jsdelivr.net/npm/canvas-confetti@1.9.3/dist/confetti.browser.min.js';
      s.onload = fire; document.head.appendChild(s);
    } else { fire(); }
  }
}
"""

CSS = """
.gradio-container{background:linear-gradient(160deg,#f3fdf6 0%,#eef6ff 55%,#fff7ea 100%) !important;}
#wc-head{text-align:center !important;max-width:980px;margin:0 auto !important;}
#wc-head h1{text-align:center !important;font-size:2.7rem !important;font-weight:900 !important;
  letter-spacing:-1px;margin:6px 0 4px !important;
  background:linear-gradient(90deg,#15803d,#0891b2,#d97706);
  -webkit-background-clip:text;background-clip:text;-webkit-text-fill-color:transparent;}
#wc-head p,#wc-head span{text-align:center !important;color:#475569 !important;}
button.primary,.primary>button{background:linear-gradient(90deg,#16a34a,#0ea5e9) !important;
  border:none !important;color:#fff !important;font-weight:700 !important;}
.block,.gr-box{border-radius:14px !important;}
footer{display:none !important;}
"""

_MODEL = None


def _model():
    """Lazy-load the detector (checkpoint pulled from HF, cached)."""
    global _MODEL
    if _MODEL is None:
        ckpt = hf_hub_download(HF_MODEL, "checkpoint_best_ema.pth")
        _MODEL = P.rfdetr_model(ckpt, "small")
    return _MODEL


@spaces.GPU(duration=120)
def _detect(video_path):
    """The only GPU step: RF-DETR-Seg over every frame -> tracks dict."""
    return P.detect_rfdetr(video_path, _model())


def _truth(name):
    rel = f"5/{name}.mp4"
    return {"goal": GT.get(rel), "leg": LEG.get(rel)}


def _table_html(res, truth):
    def row(metric, pred, actual, match=None):
        key = pred if match is None else match           # color on verdict, not on display text
        ok = (actual is None) or (str(key).upper() == str(actual).upper())
        col = "#16a34a" if ok else "#dc2626"
        act = f"<td style='padding:6px 14px;color:#555'>{actual}</td>" if actual is not None else "<td style='padding:6px 14px;color:#999'>—</td>"
        return (f"<tr><td style='padding:6px 14px;font-weight:600'>{metric}</td>"
                f"<td style='padding:6px 14px;color:{col};font-weight:700'>{pred}</td>{act}</tr>")
    goal_p = res.goal.upper() + (f" @ {res.goal_time:.1f}s" if res.goal == "goal" and res.goal_time else "")
    rows = (row("⚽ Goal", goal_p, (truth.get("goal") or "").upper() or None, match=res.goal)
            + row("🦶 Shooting foot", res.foot, truth.get("leg")))
    return ("<div style='font-weight:600;color:#666;margin-bottom:6px'>Result</div>"
            "<table style='border-collapse:collapse;font-size:15px;width:100%'>"
            "<tr style='border-bottom:2px solid #ddd;color:#999;font-size:11px'>"
            "<th style='text-align:left;padding:4px 14px'>METRIC</th>"
            "<th style='text-align:left;padding:4px 14px'>PREDICTED</th>"
            "<th style='text-align:left;padding:4px 14px'>ACTUAL</th></tr>"
            f"{rows}</table>")


def _speed_html(res):
    spd = f"{res.speed_kmh:.0f}" if res.speed_kmh else "—"
    return ("<div style='text-align:center;padding:10px'>"
            "<div style='font-weight:600;color:#666'>💨 Shot speed</div>"
            f"<div style='font-size:40px;font-weight:800;color:#0ea5e9;line-height:1.1'>{spd}</div>"
            "<div style='color:#999;font-size:13px'>km/h (ball)</div></div>")


def _coach_loading_html():
    return ("<div style='padding:14px;border:1px dashed #cbd5e1;border-radius:10px;color:#64748b'>"
            "🧠 <b>AI Coach</b> — reading your technique… "
            "<i>(Nemotron-VL via Modal)</i></div>")


def _coach_html(text):
    body = text.strip().replace("\n", "<br>")
    return ("<div style='padding:16px;border:1px solid #bbf7d0;background:#f0fdf4;border-radius:10px'>"
            "<div style='font-weight:700;color:#065f46;margin-bottom:8px'>🧠 AI Coach "
            "<span style='font-weight:400;color:#94a3b8;font-size:12px'>· Nemotron-Nano-VL (NVIDIA)</span></div>"
            f"<div style='font-size:14px;color:#334155;line-height:1.55'>{body}</div></div>")


def _coach_error_html():
    return ("<div style='padding:16px;border:1px solid #fed7aa;background:#fff7ed;border-radius:10px;color:#9a3412'>"
            "⚠️ <b>Couldn't reach the AI coach.</b> Your shot analysis above is complete — "
            "the coaching add-on is offline right now.</div>")


def analyse_clip(video_path):
    if not video_path:
        yield None, "", "", None, "", ""
        return
    name = os.path.splitext(os.path.basename(video_path))[0]
    truth = _truth(name)
    out_path = f"{RENDER_DIR}/{name}.mp4"; pose_path = f"{RENDER_DIR}/{name}_pose.png"
    tr = _detect(video_path)                                # GPU
    res = P.analyse(tr, video_path)                         # CPU
    render(video_path, tr, res, out_path)                  # CPU
    pose = pose_card_img(video_path, tr, res, pose_path)   # CPU
    table, speed = _table_html(res, truth), _speed_html(res)
    # PHASE 1 — core RF-DETR results show immediately
    yield out_path, table, speed, pose, res.goal, _coach_loading_html()
    # PHASE 2 — Nemotron-VL coach via Modal (graceful failure)
    try:
        text = C.get_coaching(C.build_payload(res), C.key_frames_b64(video_path, res))
        coach_html = _coach_html(text)
    except C.CoachUnavailable:
        coach_html = _coach_error_html()
    yield out_path, table, speed, pose, res.goal, coach_html


with gr.Blocks(title="FUT World Cup Coach", css=CSS) as demo:
    gr.Markdown(
        "# ⚽ FUT World Cup Coach 🏆\n\n"
        "Upload a football shot clip and it detects goal/no-goal (+ timestamp), who shot it, "
        "their shooting foot, and shot speed using one small RF-DETR model, plus an NVIDIA "
        "Nemotron vision-AI coach that grades your technique and tells you how to fix it. "
        "All open, local models: the detector was distilled from SAM3 + NVIDIA LocateAnything "
        "with zero manual labels. Built for the Build-Small Hackathon (World Cup 2026).",
        elem_id="wc-head")
    with gr.Row():
        with gr.Column(scale=0, min_width=42):
            gr.Markdown(EDGE_STRIP)
        with gr.Column(scale=1):
            inp = gr.Video(label="Upload a shot clip", sources=["upload"], height=300)
            btn = gr.Button("▶  Analyse", variant="primary")
            gr.Markdown("#### Examples (held-out test clips)")
            gr.Examples(examples=[[p] for p in EXAMPLES], inputs=inp, examples_per_page=12)
        with gr.Column(scale=2):
            out = gr.Video(label="Analysis", autoplay=True, height=420)
            with gr.Row():
                table = gr.HTML()
                speed = gr.HTML()
                pose = gr.Image(label="Pose Detection", height=240)
            coach = gr.HTML()                          # AI Coach card (fills in after core results)
        with gr.Column(scale=0, min_width=42):
            gr.Markdown(EDGE_STRIP)
    goal_flag = gr.Textbox(visible=False)              # carries the goal verdict to the confetti js
    btn.click(analyse_clip, inp, [out, table, speed, pose, goal_flag, coach]).then(
        None, goal_flag, None, js=CONFETTI_JS)
    inp.change(analyse_clip, inp, [out, table, speed, pose, goal_flag, coach]).then(
        None, goal_flag, None, js=CONFETTI_JS)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="amber"))