"""hf-pipe console — minimal embed UI, a projection of the spec.

The demo point: nothing here is hand-modeled. The form fields come from
TaskSpec.model_json_schema() (types/enums/defaults), the model dropdown from
the catalogue (tested paths + receipts), the prefill from resolve() — the
same three surfaces the CLI uses. One task (embed), one column of UI.
"""

from __future__ import annotations

import json

import gradio as gr
from jobpipe import TaskSpec, compile
from jobpipe.catalogue import entries
from jobpipe.resolve import resolve

SCHEMA = TaskSpec.model_json_schema()
FLAVORS = SCHEMA["properties"]["flavor"]["enum"]  # schema-derived, not hand-listed
CURATED = [e["model"] for e in entries("embeddings")]
ANY = "(any model — unverified path)"


def _tok(oauth_token, pasted):
    if oauth_token is not None and getattr(oauth_token, "token", None):
        return oauth_token.token
    return pasted or None


def _resolve(dataset, model, column, flavor, limit, world, output, config, split, engine, batch, token, oauth_token: gr.OAuthToken | None = None):
    token = _tok(oauth_token, token)
    if not dataset.strip():
        return "enter a dataset id", "", gr.update(interactive=False)
    kw = {}
    if model and model != CURATED[0]:
        kw["model"] = None if model == ANY else model
    if column.strip():
        kw["column"] = column.strip()
    if flavor:
        kw["flavor"] = flavor
    if limit:
        kw["limit"] = int(limit)
    if world and int(world) > 1:
        kw["world"] = int(world)
    for name, val in (("output", output), ("config", config), ("split", split),
                      ("engine", engine)):
        if val and str(val).strip():
            kw[name] = str(val).strip()
    if batch:
        kw["batch"] = int(batch)
    try:
        r = resolve("embeddings", dataset.strip(), token=token or None,
                    **{k: v for k, v in kw.items() if v is not None})
        c = compile(r.spec, token=token or None)
    except Exception as e:  # surfaced, not swallowed — the honest failure mode
        return f"resolution failed: {e}", "", gr.update(interactive=False)
    prov = "\n".join(
        f"{f:<8} {getattr(r.spec, f)!r}  [{r.provenance[f]}]"
        + (f"  — {r.notes[f]}" if f in r.notes else "")
        for f in ("dataset", "config", "split", "column", "model", "engine",
                  "flavor", "world", "output", "limit") if f in r.provenance
    )
    spec_json = json.dumps(r.spec.model_dump(), indent=2)
    return prov, f"```json\n{spec_json}\n```\n\n**driver** (sha256 `{c.run_json['driver']['sha256'][:12]}…`):\n```python\n{c.driver}\n```", gr.update(interactive=True)


def _launch(dataset, model, column, flavor, limit, world, output, config, split, engine, batch, token, oauth_token: gr.OAuthToken | None = None):
    token = _tok(oauth_token, token)
    if not token:
        return "sign in with HF (or paste a token) to launch — resolve is free", gr.update()
    from jobpipe import embed

    kw = dict(model=None if model in (ANY, "") else model,
              column=column.strip() or None, flavor=flavor or None,
              limit=int(limit) if limit else None,
              world=int(world) if world and int(world) > 1 else None,
              output=(output or "").strip() or None,
              config=(config or "").strip() or None,
              split=(split or "").strip() or None,
              engine=engine or None,
              batch=int(batch) if batch else None,
              token=token, verbose=False)
    try:
        run = embed(dataset.strip(), **{k: v for k, v in kw.items() if v is not None})
    except Exception as e:
        return f"launch failed: {e}", gr.update()
    repo = "/".join(run.output.removeprefix("hf://datasets/").split("/")[:2])
    url = f"https://huggingface.co/datasets/{repo}"
    return (f"run **{run.run_id}** — {len(run.jobs)} job(s) launched\n\n"
            f"output: [{run.output}]({url})\n\n"
            f"watch: `hf pipe status '{run.output}'`"), run.output


def _status(output, token, oauth_token: gr.OAuthToken | None = None):
    token = _tok(oauth_token, token)
    if not output or not str(output).strip():
        return "no run yet — launch one, or paste an output URI above"
    from jobpipe.status import render_human, status

    try:
        doc = status(str(output).strip(), include_jobs=bool(token), token=token)
    except Exception as e:
        return f"status failed: {e}"
    return render_human(doc)


THEME = gr.themes.Monochrome(
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
)
CSS = """
.gradio-container {max-width: 880px !important; margin: 0 auto}
#prov textarea, #prov .cm-editor {font-size: 12.5px}
footer {display: none !important}
"""

with gr.Blocks(title="hf pipe — embed any dataset", theme=THEME, css=CSS) as demo:
    gr.Markdown("## embed any dataset\n"
                "One launch on HF Jobs. The form is generated from the task "
                "schema, the model list is the receipted catalogue, and "
                "**resolve** shows every decision (and where it came from) "
                "before anything spends. Same seams as `hf pipe embed`.")
    with gr.Row():
        dataset = gr.Textbox(label="dataset", placeholder="fka/prompts.chat", scale=3)
        column = gr.Textbox(label="column (blank = detect)", scale=2)
    with gr.Row():
        model = gr.Dropdown(label="model (catalogue = tested paths)",
                            choices=[*CURATED, ANY], value=CURATED[0], scale=3,
                            allow_custom_value=True)
        flavor = gr.Dropdown(label="flavor", choices=[None, *FLAVORS], scale=2)
        limit = gr.Number(label="limit (blank = all)", precision=0, scale=1)
        world = gr.Number(label="jobs (fan-out)", precision=0, value=1, minimum=1, scale=1)
    with gr.Accordion("advanced (blank = resolved for you)", open=False):
        with gr.Row():
            output = gr.Textbox(label="output URI (hf://datasets/… or hf://buckets/…)", scale=3)
            engine = gr.Dropdown(label="engine", choices=[None, "tei", "vllm"], scale=1)
        with gr.Row():
            config = gr.Textbox(label="config", scale=1)
            split = gr.Textbox(label="split", scale=1)
            batch = gr.Number(label="batch (texts/request)", precision=0, scale=1)
    gr.LoginButton()
    token = gr.Textbox(label="…or paste a token (local use; write + jobs)",
                       type="password")
    with gr.Row():
        resolve_btn = gr.Button("resolve (free)")
        launch_btn = gr.Button("launch on Jobs", variant="primary", interactive=False)
    prov = gr.Code(label="resolved spec — where every value came from",
                   language=None, elem_id="prov")
    with gr.Accordion("the exact code that will run (driver + spec)", open=False):
        detail = gr.Markdown()
    result = gr.Markdown()
    with gr.Group():
        with gr.Row():
            watch_uri = gr.Textbox(label="run view — output URI (filled by launch)", scale=4)
            refresh = gr.Button("refresh", scale=1)
        status_box = gr.Code(label="progress (read from storage; auto-refreshes every 30s)",
                             language=None)
        poll = gr.Timer(30)
    resolve_btn.click(_resolve, [dataset, model, column, flavor, limit, world, output, config, split, engine, batch, token],
                      [prov, detail, launch_btn], api_name="resolve")
    launch_btn.click(_launch, [dataset, model, column, flavor, limit, world, output, config, split, engine, batch, token],
                     [result, watch_uri], api_name="launch")
    refresh.click(_status, [watch_uri, token], [status_box], api_name="status")
    poll.tick(_status, [watch_uri, token], [status_box])

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