"""Embedding-fleet control plane — run-level view of a fan-out embedding run.

Renders one run of launch-embedding-fleet.py (uv-scripts/embeddings): documents-processed
progress, tokens, live ~$ cost vs the hard ceiling, ETA, GPU utilization, replica health,
and a per-worker table. Polls the run bucket + Jobs API every few seconds.
"""

import os

import gradio as gr

from control_plane import RunView, list_runs, load_run

BUCKET = os.environ.get("FLEET_BUCKET", "davanstrien/embedding-runs")
NAMESPACE = os.environ.get("FLEET_NAMESPACE") or BUCKET.split("/")[0]
POLL_SECS = float(os.environ.get("FLEET_POLL_SECS", "5"))

CSS = """
:root {
  --ink: #1a1a1a; --ink-2: #555; --ink-3: #999;
  --surface: #fffef9; --card: #ffffff; --line: #e4e2da;
  --accent: #3d6ea5; --ok: #2e7d43; --bad: #b3382c;
}
@media (prefers-color-scheme: dark) {
  :root { --ink: #ececec; --ink-2: #b0b0b0; --ink-3: #7a7a7a;
          --surface: #131313; --card: #1c1c1c; --line: #333;
          --accent: #7ba7d4; --ok: #6fbf85; --bad: #e07a6e; }
}
.gradio-container { max-width: 1080px !important; }
#fleet-html h2 { font-weight: 600; margin: 0 0 2px; }
.fleet-head { color: var(--ink-2); font-size: 0.92rem; margin-bottom: 14px; }
.fleet-head a { color: var(--accent); text-decoration: none; }
.tiles { display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));
         gap: 10px; margin: 12px 0 6px; }
.tile { background: var(--card); border: 1px solid var(--line); border-radius: 6px;
        padding: 10px 14px; }
.tile .k { font-size: 0.72rem; text-transform: uppercase; letter-spacing: 0.05em;
           color: var(--ink-3); }
.tile .v { font-size: 1.45rem; font-variant-numeric: tabular-nums; color: var(--ink); }
.tile .s { font-size: 0.78rem; color: var(--ink-2); }
.tile .v.ok { color: var(--ok); } .tile .v.bad { color: var(--bad); }
.bar-wrap { margin: 10px 0 2px; }
.bar-label { display: flex; justify-content: space-between; font-size: 0.85rem;
             color: var(--ink-2); margin-bottom: 4px; font-variant-numeric: tabular-nums; }
.bar { height: 10px; background: var(--line); border-radius: 5px; overflow: hidden; }
.bar > div { height: 100%; background: var(--accent); border-radius: 5px 0 0 5px;
             transition: width 0.6s ease; }
"""


def fmt_int(n):
    return f"{n:,}" if n is not None else "—"


def fmt_secs(s):
    if s is None:
        return "—"
    if s < 90:
        return f"{s:.0f}s"
    if s < 5400:
        return f"{s / 60:.0f} min"
    return f"{s / 3600:.1f} h"


def render_run(view: RunView) -> str:
    m = view.manifest
    pct = 100.0 * view.rows_done / view.rows_total if view.rows_total else 0.0
    in_url = f"https://huggingface.co/datasets/{m['input_dataset']}"
    out_url = f"https://huggingface.co/datasets/{m['output_dataset']}"
    ceiling = f"ceiling ≤ ${view.cost_ceiling_usd:,.2f}" if view.cost_ceiling_usd else "est."
    gpu = f"{view.gpu_util:.0f}%" if view.gpu_util is not None else "—"
    health_cls = "bad" if view.errored else "ok"
    eta = "done" if view.eta_secs == 0 else fmt_secs(view.eta_secs)
    return f"""
<h2>{m["output_dataset"].split("/")[-1]}</h2>
<div class="fleet-head">
  <a href="{in_url}">{m["input_dataset"]}</a> → <a href="{out_url}">{m["output_dataset"]}</a>
  · <code>{m["model"].split("/")[-1]}</code>
  · {view.num_shards} × {m["flavor"]} · run <code>{view.run_id}</code>
</div>
<div class="bar-wrap">
  <div class="bar-label">
    <span>{fmt_int(view.rows_done)} of {fmt_int(view.rows_total)} documents</span>
    <span>{pct:.1f}%</span>
  </div>
  <div class="bar"><div style="width:{min(pct, 100):.2f}%"></div></div>
</div>
<div class="tiles">
  <div class="tile"><div class="k">Tokens (est.)</div>
    <div class="v">{fmt_int(view.tokens_done_est)}</div></div>
  <div class="tile"><div class="k">Cost</div>
    <div class="v">~${view.cost_usd:,.2f}</div><div class="s">{ceiling}</div></div>
  <div class="tile"><div class="k">ETA</div>
    <div class="v">{eta}</div></div>
  <div class="tile"><div class="k">GPU util</div>
    <div class="v">{gpu}</div><div class="s">mean of running</div></div>
  <div class="tile"><div class="k">Replicas healthy</div>
    <div class="v {health_cls}">{view.healthy}/{view.num_shards}</div>
    <div class="s">{view.done} done · {view.errored} error</div></div>
</div>
"""


def worker_table(view: RunView):
    rows = []
    for w in view.workers:
        rows.append([
            w.rank,
            w.stage + (f" ({w.state})" if w.state and w.state != "running" else ""),
            f"{w.rows_done:,}" + (f" / {w.rows_total:,}" if w.rows_total else ""),
            f"{w.rows_per_sec:,.0f}" if w.rows_per_sec else "—",
            f"{w.gpu_util:.0f}%" if w.gpu_util is not None else "—",
            f"~${w.cost_usd:.3f}" if w.cost_usd is not None else "—",
            w.job_id or "—",
        ])
    return rows


def refresh(run_id):
    if not run_id:
        return "<p>No runs found in the bucket yet.</p>", []
    view = load_run(BUCKET, run_id, namespace=NAMESPACE)
    if view is None:
        return f"<p>Run <code>{run_id}</code> has no manifest.</p>", []
    return render_run(view), worker_table(view)


def init(request: gr.Request):
    runs = list_runs(BUCKET)
    wanted = dict(request.query_params).get("run")
    selected = wanted if wanted in runs else (runs[0] if runs else None)
    html, table = refresh(selected)
    return gr.Dropdown(choices=runs, value=selected), html, table


with gr.Blocks(css=CSS, title="Embedding Fleet") as demo:
    with gr.Row():
        # allow_custom_value: ?run= deep links and API calls may reference runs
        # that appeared after the choices list was built.
        run_dd = gr.Dropdown(label="Run", choices=[], scale=3, allow_custom_value=True)
        reload_btn = gr.Button("Reload runs", scale=1)
    html = gr.HTML(elem_id="fleet-html")
    table = gr.Dataframe(
        headers=["rank", "stage", "rows", "rows/s", "gpu", "~$", "job"],
        interactive=False, label="Workers",
    )
    timer = gr.Timer(POLL_SECS)

    demo.load(init, inputs=None, outputs=[run_dd, html, table])
    timer.tick(refresh, inputs=run_dd, outputs=[html, table])
    run_dd.change(refresh, inputs=run_dd, outputs=[html, table])
    reload_btn.click(lambda: gr.Dropdown(choices=list_runs(BUCKET)), outputs=run_dd)

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