"""HAWK ONE - watch a play run, redraw a player's path, see what would have happened.   Run:  python app.py"""
import json
import math
import os
import random
import tempfile

import gradio as gr

import tracking
import whatif
from coach import DefCoach, OffCoach, fourth_down, snap_spot
from film import build_film_html
from playbook import DEFENSE_CALLS, OFFENSE_PLAYS
from render import CSS as FILM_CSS, lab_html, outcome_html
from rules import PROFILES

import spaces
# This dummy function satisfies the Hugging Face ZeroGPU startup validator
@spaces.GPU(duration=1)
def dummy_gpu_initializer():
    return "GPU Initialized"

LEVELS = [("NFL", "NFL"), ("College (NCAA)", "NCAA")]
SPOTS = ["Left hash", "Middle", "Right hash"]
COACH = "Coach decides"

HEADER = """
<style>
.gradio-container{max-width:1500px!important}
.hk-top{display:flex;align-items:center;gap:14px;margin:2px 0 6px}
.hk-top .logo{font:800 30px 'Barlow Condensed','Arial Narrow',Impact,sans-serif;letter-spacing:1px;color:#fff;background:#d50a0a;padding:2px 16px 2px 12px;clip-path:polygon(0 0,100% 0,calc(100% - 12px) 100%,0 100%)}
.hk-top .tag{color:#64748b;font-size:13px}
.hk-info{font:500 13.5px Inter,system-ui,sans-serif;color:#cfdcf5;background:#0a1b38;border:1px solid #1f3a6b;border-left:4px solid #d50a0a;border-radius:6px;padding:8px 14px;line-height:1.5}
.hk-info b{color:#fff}.hk-info .bn{display:inline-block;background:#d50a0a;color:#fff;font-weight:800;letter-spacing:.8px;padding:0 9px;border-radius:3px;margin-right:8px}
</style>
<div class="hk-top"><span class="logo">HAWK ONE</span><span class="tag">Film room &middot; one engine, NFL and college rules</span></div>
"""

BRIDGE_JS = """() => {
  if (window.__hawkBridge) return; window.__hawkBridge = true;
  window.addEventListener('message', function (e) {
    var d = e.data; if (!d || !d.hawk) return;
    var box = document.querySelector('#hawk_bridge textarea, #hawk_bridge input, textarea#hawk_bridge, input#hawk_bridge'); if (!box) return;
    var proto = box.tagName === 'TEXTAREA' ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
    Object.getOwnPropertyDescriptor(proto, 'value').set.call(box, JSON.stringify({kind: d.kind, edits: d.edits, react: d.react, sel: d.sel}));
    box.dispatchEvent(new Event('input', {bubbles: true}));
    if (d.kind === 'run') setTimeout(function () {
      var b = document.querySelector('#hawk_run'); if (b && b.tagName !== 'BUTTON') b = b.querySelector('button');
      if (b) b.click();
    }, 300);
  });
}"""

EMPTY = ("<div style='padding:60px;text-align:center;color:#64748b;border:1px dashed #94a3b8;border-radius:10px'>"
         "Choose the play calls (or let the coaches decide) and press <b>Run play</b>.</div>")


# ---------------------------------------------------------------- helpers
def _ydl(yl):
    yl = int(round(yl))
    return f"own {yl}" if yl <= 50 else f"opp {100 - yl}"


def _pick(off_play, def_call, sit, seed):
    """Coach calls fill in anything left on 'Coach decides'; returns calls plus the reason for each."""
    rng = random.Random(int(seed) * 7919 + 13)
    s = dict(down=sit["down"], dist=sit["dist"], yl=sit["yl"])
    op, ow = OffCoach(rng).call(s) if off_play == COACH else (off_play, "your call")
    dc, dw = DefCoach(rng).call(s) if def_call == COACH else (def_call, "your call")
    return op, dc, ow, dw


def _info(st, banner=""):
    s = st["sit"]
    b = f"<span class='bn'>{banner}</span>" if banner else ""
    return (f"<div class='hk-info'>{b}<b>{['1st', '2nd', '3rd', '4th'][s['down'] - 1]} &amp; {s['dist']}</b> at {_ydl(s['yl'])} &middot; "
            f"Offense <b>{st['setup']['off_play']}</b> <span style='opacity:.7'>({st['why'][0]})</span> &middot; "
            f"Defense <b>{st['setup']['def_call']}</b> <span style='opacity:.7'>({st['why'][1]})</span></div>")


def _viewer(st, rep_a, rep_b=None, edits=None, sel=None, mode=None, react=None):
    s = st["sit"]
    return build_film_html(st["level"], st["setup"], rep_a, rep_b, st["paths"], edits=edits, sel=sel, mode=mode,
                           down=s["down"], dist=s["dist"], autoplay=True, react=react, source=st.get("source", "simulated"))


def _msg(bridge):
    try:
        d = json.loads(bridge) if bridge else {}
        return d.get("edits") or {}, d.get("sel"), d.get("react") or {}
    except Exception:
        return {}, None, {}


def _start(level, off_play, def_call, o_rat, d_rat, team_seed, sit, play_no, banner=""):
    op, dc, ow, dw = _pick(off_play, def_call, sit, play_no)
    setup = whatif.make_setup(level, op, dc, o_rat, d_rat, team_seed, int(sit["yl"]), by=sit["by"])
    res, rep = whatif.simulate(setup, int(play_no), None, True)
    st = dict(level=level, setup=setup, paths=whatif.extract_paths(rep), rep=rep, play_no=int(play_no), sit=dict(sit),
              why=(ow, dw), res=res, lab=None, args=(o_rat, d_rat, team_seed))
    chs = [(f"{m['pos']} #{m['num']}  ({m['id']})", m["id"]) for m in rep["meta"]]
    first = (rep.get("marks") or {}).get("target") or next(
        (m["id"] for m in rep["meta"] if m["pos"] in ("WR", "RB", "TE")), rep["meta"][0]["id"])
    return (st, _viewer(st, rep, sel=first), "", "", gr.update(choices=chs, value=first), "",
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), _info(st, banner))


def run_play(level, off_play, def_call, o_rat, d_rat, team_seed, yl, spot, play_no, down, dist):
    pf = PROFILES[level]
    by = whatif.make_setup(level, "Quick Slants", "4-3 Over · Cover 3", spot=spot)["by"]
    return _start(level, off_play, def_call, o_rat, d_rat, team_seed,
                  dict(yl=int(yl), by=by, down=int(down), dist=int(dist)), int(play_no))


def next_play(st, level, off_play, def_call, o_rat, d_rat, team_seed, yl, spot, play_no, down, dist):
    """Advance the drive from the last result and run the next snap."""
    if not st or st["level"] != level or st.get("source") == "recorded":
        return run_play(level, off_play, def_call, o_rat, d_rat, team_seed, yl, spot, play_no, down, dist)
    pf, sit, res = PROFILES[level], dict(st["sit"]), st["res"]
    banner, fresh = "", False
    if res["td"]:
        banner, fresh = "TOUCHDOWN - new drive", True
    elif res["turnover"] or res["outcome"] == "interception":
        banner, fresh = "TURNOVER - new drive", True
    elif res["penalty"]:
        pen = res["penalty"]
        banner = f"FLAG - {pen['kind']}"
        if pen["kind"] == "dpi":
            sit["yl"] = min(99, sit["yl"] + int(round(pen["yards"])))
            sit["down"], sit["dist"] = 1, min(10, 100 - sit["yl"])
        else:
            sit["yl"] = max(1, sit["yl"] + pen["yards"])
            sit["dist"] = min(99, sit["dist"] - pen["yards"])
    else:
        gain = res["yards"]
        sit["yl"] = max(1, min(99, sit["yl"] + int(round(gain))))
        if res["outcome"] != "incomplete":
            sit["by"] = snap_spot(pf, res["end_y"])
        if gain >= sit["dist"]:
            sit["down"], sit["dist"] = 1, min(10, 100 - sit["yl"])
        else:
            sit["down"] += 1
            sit["dist"] = max(1, int(math.ceil(sit["dist"] - gain)))
            if sit["down"] > 4:
                banner, fresh = "TURNOVER ON DOWNS - new drive", True
    if not fresh and sit["down"] == 4:
        dec = fourth_down(sit, random.Random(st["play_no"]))
        if dec != "go":
            banner, fresh = ("FIELD GOAL attempt" if dec == "fg" else "PUNT") + " - new drive", True
    if fresh:
        sit = dict(yl=25, by=whatif.make_setup(level, "Quick Slants", "4-3 Over · Cover 3", spot="Middle")["by"], down=1, dist=10)
    o, d, ts = st["args"]
    return _start(level, off_play, def_call, o, d, ts, sit, st["play_no"] + 1, banner)


def open_recorded(st, level, file):
    """Load a play from a tracking CSV (see tracking.py for the format)."""
    keep = [gr.update()] * 7                      # lab, focus, bridge, three apply buttons, info: leave as they are
    if not file:
        return (st, gr.update(), "<p>Choose a tracking CSV first.</p>", *keep)
    try:
        path = file if isinstance(file, str) else getattr(file, "name", file)
        text = open(path, encoding="utf-8-sig").read()
        pid = tracking.plays_in(text)[0]
        rep = tracking.load_csv(text, pid)
    except Exception as e:
        return (st, gr.update(), f"<p>Could not read that file: {e}</p>", *keep)
    yl = int(round(rep["bx"] - 10))
    setup = dict(level=level, off_play="Recorded play", def_call=f"play {pid}", bx=rep["bx"], by=rep["by"], o_rat=0, d_rat=0, team_seed=0)
    st = dict(level=level, setup=setup, paths=whatif.extract_paths(rep), rep=rep, play_no=0,
              sit=dict(yl=yl, by=rep["by"], down=1, dist=10), why=("recorded", "recorded"), res=None, lab=None,
              args=(78, 75, 7), source="recorded")
    chs = [(f"{m['pos']} #{m['num']}  ({m['id']})", m["id"]) for m in rep["meta"]]
    first = rep["marks"].get("target") or rep["meta"][0]["id"]
    info = (f"<div class='hk-info'><span class='bn'>RECORDED</span>Play {pid} from your tracking file &middot; "
            f"{len(rep['meta'])} players &middot; inspect separation and speed, draw a suggested path. "
            f"What-if runs on simulated plays.</div>")
    return (st, _viewer(st, rep, sel=first), "", "", gr.update(choices=chs, value=first), "",
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), info)


def export_csv(st):
    if not st:
        return None
    path = os.path.join(tempfile.mkdtemp(), "hawk_one_play.csv")
    with open(path, "w", encoding="utf-8") as f:
        f.write(tracking.to_csv(st["rep"], st.get("play_no", 1) or 1))
    return path


def restart(st):
    if not st:
        return st, EMPTY, ""
    return st, _viewer(st, st["rep"]), ""


def film_run(st, bridge, n, level):
    if not st:
        return st, EMPTY, "<p>Run a play first.</p>"
    if st["level"] != level:
        return st, _viewer(st, st["rep"]), "<p>The level changed after this play ran. Press <b>Run play</b> again.</p>"
    if st.get("source") == "recorded":
        return st, _viewer(st, st["rep"]), ("<p>What-if needs a simulated play: the engine must know the play design and the coverage to let "
                                            "the other players react. Recorded plays can be inspected and annotated here; running them through the "
                                            "engine is the next step once real data is available.</p>")
    edits, sel, react = _msg(bridge)
    edits = whatif.changed(edits, st["paths"])
    react = {k: v for k, v in react.items() if abs(v) > 1e-6}
    if not edits and not react:
        return st, _viewer(st, st["rep"], sel=sel), "<p>Nothing has been changed yet. Select a player, drag a point on his path or change his reaction time, then run.</p>"
    return (st, *_compare_and_show(st, edits, sel, n, react))


def _compare_and_show(st, edits, sel, n, react=None):
    react = react or {}
    pool = list(edits) + list(react)
    focus = sel if sel in pool else pool[0]
    c = whatif.compare(st["setup"], edits, st["paths"], n=int(n), seed0=st["play_no"], focus=focus,
                       need=st["sit"]["dist"], react=react)
    st["cmp"] = c
    return _viewer(st, c["rep_base"], c["rep_alt"], c["scripts"], focus, "overlay", c["react"]), outcome_html(c, st["sit"]["dist"])


def film_lab(st, bridge, focus, n):
    none = [gr.update(visible=False)] * 3
    if not st:
        return st, "<p>Run a play first.</p>", *none
    _, sel, _r = _msg(bridge)
    pid = sel or focus
    m = next((m for m in st["rep"]["meta"] if m["id"] == pid), None)
    if m is None:
        return st, "<p>Pick a player.</p>", *none
    sg = whatif.suggest(st["setup"], pid, st["paths"], n=max(40, int(n) // 4), seed0=st["play_no"], need=st["sit"]["dist"])
    st["lab"] = dict(pid=pid, rows=sg["rows"])
    btns = [gr.update(visible=i < len(sg["rows"]), value=f"Apply #{i + 1}: {sg['rows'][i]['name']}" if i < len(sg["rows"]) else "")
            for i in range(3)]
    return st, lab_html(pid, f"{m['pos']} #{m['num']}", sg), *btns


def film_apply(k):
    def go(st, n):
        if not st or not st.get("lab") or k >= len(st["lab"]["rows"]):
            return st, EMPTY, ""
        lab = st["lab"]
        return (st, *_compare_and_show(st, {lab["pid"]: lab["rows"][k]["path"]}, lab["pid"], n))
    return go


# ---------------------------------------------------------------- UI
with gr.Blocks(title="HAWK ONE") as demo:
    gr.HTML(HEADER + FILM_CSS)
    with gr.Row(equal_height=True):
        level = gr.Radio(LEVELS, value="NFL", label="Level of play", scale=2)
        f_off = gr.Dropdown([COACH] + list(OFFENSE_PLAYS), value=COACH, label="Offense", scale=2)
        f_def = gr.Dropdown([COACH] + list(DEFENSE_CALLS), value=COACH, label="Defense", scale=3)
        b_run = gr.Button("Run play", variant="primary", scale=1)
        b_next = gr.Button("Next play", scale=1)
    with gr.Accordion("Situation and teams", open=False):
        with gr.Row():
            f_down = gr.Radio([1, 2, 3, 4], value=1, label="Down")
            f_dist = gr.Slider(1, 20, value=10, step=1, label="Yards to go")
            f_yl = gr.Slider(1, 99, value=30, step=1, label="Ball on (from own goal)")
            f_spot = gr.Radio(SPOTS, value="Middle", label="Ball spot")
        with gr.Row():
            f_play = gr.Number(3, precision=0, label="Play # (scenario)")
            f_orat = gr.Slider(55, 95, value=78, step=1, label="Offense rating")
            f_drat = gr.Slider(55, 95, value=75, step=1, label="Defense rating")
            f_tseed = gr.Number(7, precision=0, label="Roster seed")
    f_info = gr.HTML()
    f_view = gr.HTML(EMPTY)
    with gr.Row():
        f_n = gr.Slider(60, 600, value=200, step=20, label="Runs per side", scale=3)
        f_run = gr.Button("Run what-if", variant="primary", elem_id="hawk_run", scale=1)
        f_clear = gr.Button("Replay original", scale=1)
    f_out = gr.HTML()
    with gr.Accordion("Path Lab - let the engine try variations of one player's path", open=False):
        with gr.Row():
            f_focus = gr.Dropdown([], label="Player (or the one selected in the viewer)", scale=3)
            f_labbtn = gr.Button("Find better paths", scale=1)
        f_lab = gr.HTML()
        with gr.Row():
            f_a1 = gr.Button("Apply #1", visible=False)
            f_a2 = gr.Button("Apply #2", visible=False)
            f_a3 = gr.Button("Apply #3", visible=False)
    with gr.Accordion("Recorded plays (tracking data)", open=False):
        gr.Markdown("Load a play from a tracking CSV (columns: `play_id, t, entity, side, pos, num, x, y, z`; format notes in `tracking.py`). "
                    "Use **Save this play as CSV** to get a sample file in exactly that format.")
        with gr.Row():
            f_file = gr.File(label="Tracking CSV", file_types=[".csv"], type="filepath", scale=3)
            f_open = gr.Button("Open recorded play", scale=1)
            f_save = gr.Button("Save this play as CSV", scale=1)
            f_csv = gr.File(label="Download", scale=2)
    f_bridge = gr.Textbox(elem_id="hawk_bridge", label="bridge", lines=1)
    film_state = gr.State(None)

    demo.load(None, None, None, js=BRIDGE_JS)
    sit_in = [level, f_off, f_def, f_orat, f_drat, f_tseed, f_yl, f_spot, f_play, f_down, f_dist]
    outs = [film_state, f_view, f_out, f_lab, f_focus, f_bridge, f_a1, f_a2, f_a3, f_info]
    b_run.click(run_play, sit_in, outs)
    level.change(run_play, sit_in, outs)
    b_next.click(next_play, [film_state] + sit_in, outs)
    f_open.click(open_recorded, [film_state, level, f_file], outs)
    f_save.click(export_csv, film_state, f_csv)
    f_run.click(film_run, [film_state, f_bridge, f_n, level], [film_state, f_view, f_out])
    f_clear.click(restart, film_state, [film_state, f_view, f_out])
    f_labbtn.click(film_lab, [film_state, f_bridge, f_focus, f_n], [film_state, f_lab, f_a1, f_a2, f_a3])
    for k, b in enumerate((f_a1, f_a2, f_a3)):
        b.click(film_apply(k), [film_state, f_n], [film_state, f_view, f_out])
    demo.load(run_play, sit_in, outs)          # the first play runs as soon as the page opens

# if __name__ == "__main__":
demo.launch()