"""LTX-2 LoRA Trainer — HF Space (HF Jobs + buckets).

Sign in with Hugging Face, upload videos + captions, set hyperparameters, and submit a
training job to HF Jobs. The job runs on a GPU flavor, reproduces the trainer env from the
lockfile, trains a LoRA / IC-LoRA for LTX-2.3, and pushes it to your Hub repo. Datasets are
staged on HF buckets. Everything runs under the *signed-in user's* account — no pasted tokens.

Runs on `cpu-basic` — the Space only submits + monitors jobs (no GPU, no torch here).
"""

from __future__ import annotations

import re

import gradio as gr

import jobs

FLAVORS = ["a100-large", "a100x4", "l40sx1", "l40sx4"]
MAX_LOG = 60_000
MODE_KEYS = list(jobs.MODES.keys())

CSS = """
#hdr {text-align:left; padding:4px 2px 0 2px;}
#hdr h1 {margin:0; font-size:1.7rem;}
#hdr p {margin:.25rem 0 0 0; color:var(--body-text-color-subdued);}
.section-card {border:1px solid var(--block-border-color); border-radius:14px;
               padding:14px 16px; background:var(--block-background-fill);}
.step-badge {font-weight:600; color:var(--primary-500);}
#banner {border-radius:12px; padding:10px 14px; font-size:.95rem;}
footer {visibility:hidden;}
"""

THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.indigo,
    secondary_hue=gr.themes.colors.purple,
    radius_size=gr.themes.sizes.radius_lg,
)


def _signin_state(profile: gr.OAuthProfile | None):
    """Banner + default hub id, recomputed on load and after sign-in."""
    if profile is None:
        return (
            "🔒 **You're not signed in.** Use **Sign in with Hugging Face** (top-right) to start — "
            "your dataset, the training job, and the resulting LoRA all run under **your** account "
            "and billing. No tokens to paste.",
            gr.update(),
        )
    return (
        f"✅ Signed in as **{profile.username}** — jobs, buckets and the pushed LoRA will live under "
        f"your account.",
        gr.update(value=f"{profile.username}/ltx2-lora"),
    )


def submit_job(
    files, captions_text, run_name, mode, resolution, rank, alpha, lr, steps, batch_size,
    grad_accum, quantization, optimizer_type, te_8bit, push, hub_id, flavor, timeout,
    profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None,
):
    if oauth_token is None or profile is None:
        return "❌ Please **sign in with Hugging Face** first (top-right).", "", ""
    if not files:
        return "❌ Upload at least one video.", "", ""
    try:
        jobs.parse_resolution(resolution)
    except ValueError as e:
        return f"❌ {e}", "", ""
    if push and not hub_id.strip():
        return "❌ Set a Hub model id (e.g. you/my-lora) or disable push.", "", ""

    params = {
        "run_name": run_name or "ltx2-lora", "mode": mode, "resolution": resolution.strip(),
        "rank": rank, "alpha": alpha, "learning_rate": lr, "steps": steps, "batch_size": batch_size,
        "gradient_accumulation_steps": grad_accum,
        "quantization": None if quantization in ("none", None) else quantization,
        "optimizer_type": optimizer_type, "load_text_encoder_in_8bit": bool(te_8bit),
        "push_to_hub": bool(push), "hub_model_id": hub_id.strip(),
        "hf_token": oauth_token.token,
        "captions": captions_text.splitlines() if captions_text else [], "seed": 42,
    }
    try:
        res = jobs.submit(params, [f for f in files], flavor=flavor, timeout=timeout)
    except Exception as e:  # noqa: BLE001
        return f"❌ Submission failed: {e}", "", ""

    status = f"✅ Job submitted on **{flavor}**, running as **{profile.username}**."
    link = ""
    if res["url"]:
        link = f"**Job:** [{res['job_id']}]({res['url']})  \n**Bucket:** `{res['bucket']}`"
    elif res["job_id"]:
        link = f"**Job id:** `{res['job_id']}`  \n**Bucket:** `{res['bucket']}`"
    return status, link, res["log"]


def refresh(job_id, oauth_token: gr.OAuthToken | None):
    if not job_id.strip():
        return "Enter a job id.", ""
    token = oauth_token.token if oauth_token else ""
    st = jobs.job_status(job_id.strip(), token)
    logs = jobs.job_logs(job_id.strip(), token)
    return f"**Status:** `{st}`", logs[-MAX_LOG:] if len(logs) > MAX_LOG else logs


def _extract_id(link_md: str) -> str:
    m = re.search(r"\[([^\]]+)\]\(http", link_md or "")
    return m.group(1) if m else ""


MODE_HELP = (
    "**IC-LoRA (in-context control)** learns from *pairs*: a target video `clip.mp4` and its "
    "control/reference video `clip_reference.mp4` (depth, pose, edges, an inpainting-masked "
    "version, …). Upload both — they're matched by filename. **T2V / I2V** need only the target clips."
)

with gr.Blocks(title="LTX-2 LoRA Trainer") as demo:
    with gr.Row(equal_height=True):
        gr.Markdown(
            "# 🎬 LTX-2 LoRA Trainer\n"
            "Train a **LoRA / IC-LoRA for LTX-2.3** on your own videos. Training runs on "
            "**HF Jobs**, data is staged on **HF buckets**, and the LoRA is pushed to your Hub — "
            "you only pay for the GPU runtime.",
            elem_id="hdr",
        )
        gr.LoginButton(scale=0, min_width=200)

    banner = gr.Markdown(elem_id="banner")

    with gr.Row():
        # ---------------------------------------------------------------- left: data + training
        with gr.Column(scale=3):
            with gr.Group():
                gr.Markdown("<span class='step-badge'>STEP 1</span> &nbsp; **Dataset**")
                mode = gr.Dropdown(MODE_KEYS, value=MODE_KEYS[0], label="Training mode")
                gr.Markdown(MODE_HELP)
                files = gr.File(
                    label="Videos — targets + (for IC-LoRA) their *_reference videos, or a .zip",
                    file_count="multiple", file_types=["video", ".zip"], height=160,
                )
                captions = gr.Textbox(
                    label="Captions — one per line, in filename order of the target clips",
                    lines=4, placeholder="a fluffy cat in a sunlit room\na sweeping green landscape\n…",
                )

            with gr.Group():
                gr.Markdown("<span class='step-badge'>STEP 2</span> &nbsp; **Training settings**")
                resolution = gr.Textbox(
                    label="Resolution (W×H×F)", value="768x512x49",
                    info="W,H divisible by 32 · frames F satisfy F % 8 == 1 (25, 49, 81, …)",
                )
                with gr.Row():
                    rank = gr.Number(label="LoRA rank", value=32, precision=0)
                    alpha = gr.Number(label="LoRA alpha", value=32, precision=0)
                with gr.Row():
                    lr = gr.Number(label="Learning rate", value=2e-4)
                    steps = gr.Number(label="Steps", value=2000, precision=0)
                with gr.Accordion("Advanced", open=False):
                    with gr.Row():
                        batch_size = gr.Number(label="Batch size", value=1, precision=0)
                        grad_accum = gr.Number(label="Grad accumulation", value=1, precision=0)
                    with gr.Row():
                        quantization = gr.Dropdown(
                            ["none", "int8-quanto", "fp8-quanto"], value="none", label="Quantization",
                            info="a100-large (80 GB) fits 22B in bf16 — quantize only on smaller GPUs.",
                        )
                        optimizer_type = gr.Dropdown(["adamw", "adamw8bit"], value="adamw", label="Optimizer")
                    te_8bit = gr.Checkbox(label="Load text encoder in 8-bit", value=False)

        # ---------------------------------------------------------------- right: output + submit + monitor
        with gr.Column(scale=2):
            with gr.Group():
                gr.Markdown("<span class='step-badge'>STEP 3</span> &nbsp; **Output & launch**")
                run_name = gr.Textbox(label="Run name", value="ltx2-ic-lora")
                push = gr.Checkbox(label="Push the trained LoRA to my Hub", value=True)
                hub_id = gr.Textbox(label="Hub model id", placeholder="username/my-lora")
                with gr.Row():
                    flavor = gr.Dropdown(FLAVORS, value=jobs.DEFAULT_FLAVOR, label="GPU flavor")
                    timeout = gr.Textbox(label="Timeout", value="4h")
                submit_btn = gr.Button("🚀 Submit training job", variant="primary", size="lg")
                status = gr.Markdown("")
                joblink = gr.Markdown("")

            with gr.Group():
                gr.Markdown("**Monitor a job**")
                with gr.Row():
                    job_id = gr.Textbox(label="Job id", scale=3)
                    refresh_btn = gr.Button("🔄 Refresh", scale=1)
                mon_status = gr.Markdown("")
                mon_logs = gr.Textbox(label="Job logs", lines=16, autoscroll=True, max_lines=16)

    sublog = gr.Textbox(label="Submission output", lines=4, visible=False)

    demo.load(_signin_state, inputs=None, outputs=[banner, hub_id])

    submit_btn.click(
        submit_job,
        inputs=[files, captions, run_name, mode, resolution, rank, alpha, lr, steps, batch_size,
                grad_accum, quantization, optimizer_type, te_8bit, push, hub_id, flavor, timeout],
        outputs=[status, joblink, sublog],
    ).then(_extract_id, inputs=joblink, outputs=job_id)
    refresh_btn.click(refresh, inputs=[job_id], outputs=[mon_status, mon_logs])


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch(theme=THEME, css=CSS)