"""Krea 2 LoRA Trainer — HF Space (HF Jobs backend).

Sign in with Hugging Face, pick what you're training (a style, or an object/character), upload a
handful of images, let the built-in AI caption them, and submit a DreamBooth-LoRA job to HF Jobs.
The job trains on **Krea 2 RAW**, (optionally) validates on **Krea 2 Turbo**, and pushes the LoRA
to your Hub — all under your account. The Space runs on `cpu-basic`; the GPU work runs on HF Jobs.

Three tokens, three jobs:
  * the **user's** OAuth token — dataset + the pushed LoRA (their account / billing);
  * `KREA_TOKEN` secret — downloads the gated Krea 2 weights inside the job only;
  * `CAPTION_HF_TOKEN` secret — calls the Inference API for AI captioning on this Space only.
"""

from __future__ import annotations

import os

import gradio as gr

import caption
import jobs

MAX_IMAGES = 40
MAX_LOG = 60_000
IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tiff", ".tif")

LR_SCHEDULERS = ["constant", "cosine", "linear", "constant_with_warmup", "polynomial"]
OPTIMIZERS = ["adamW", "prodigy"]
QUANT_CHOICES = [
    ("None — bf16 (best quality, most VRAM)", "none"),
    ("FP8 (faster compute, needs GPU ≥ 8.9)", "fp8"),
    ("4-bit NF4 / QLoRA (lowest VRAM)", "4bit"),
]
CONCEPT_CHOICES = [
    ("🎨 Style", "style"),
    ("🧸 Object or character", "object"),
    ("⚙️ Custom", "custom"),
]

AGENT_URL = "https://huggingface.co/spaces/multimodalart/krea2-lora-trainer/raw/main/agent.md"

CSS = """
#main_title{text-align:center}
#main_title h1{font-size:2.25rem;margin-bottom:0}
#main_title h3{margin-top:.25em;font-size:1.25em}
#main_title p{margin-top:.25em;font-size:1.05em;opacity:.85}
.accordion{color:var(--body-text-color)}
.login_logout{width:100% !important}
#login{width:100% !important;margin:0 auto}
#signin_banner{text-align:center}
.agent_hint{text-align:center;opacity:.7;font-size:.9em;margin:.4em 0 -.2em}
.agent_curl{max-width:380px;margin:0 auto}
#main_ui.locked{pointer-events:none;opacity:.55;filter:grayscale(.3)}

/* captioning panel */
#captioning_area{padding:0 4px 4px 16px}
#captioning_area h2{margin:0 0 8px}
.trigger_row{gap:8px;align-items:stretch}
.suggest_btn button{width:100%;height:100%;border:var(--button-border-width) solid var(--button-secondary-border-color);background:var(--button-secondary-background-fill)}
.cap_item{border:1px solid var(--border-color-primary);border-radius:10px;padding:10px;margin-bottom:12px;gap:8px}
.cap_img{flex:0 0 auto !important;max-width:180px}
.cap_img .icon-button-wrapper,.cap_img .source-selection,.cap_img .image-frame .controls{display:none !important}
.cap_img img,.cap_img .image-frame{object-fit:cover;border-radius:8px;max-height:180px}
.cap_row{align-items:stretch !important;gap:10px}
.cap_one_btn button{height:100%;min-width:46px;padding-left:0;padding-right:0}
"""

THEME = gr.themes.Monochrome(
    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px",
                             xxl="24px", xxs="9px"),
    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
)

FLAVOR_GUIDE = """**Which GPU?** You're billed per-minute of actual runtime. Krea 2 is a 12B DiT.

| Flavor | VRAM | Best for |
|---|---|---|
| `rtx-pro-6000` | 96 GB | **recommended (default)** — newest Blackwell card, bf16 with room to spare |
| `a100-large` | 80 GB | proven A100 — bf16 with offload + cached latents |

First run is slow to start: it downloads the gated Krea 2 RAW + Turbo weights before training.
"""

TRIGGER_HELP = {
    "style": (
        "Your **trigger** is a short descriptive phrase appended to every caption — e.g. "
        "*heavy impasto style*, *monochrome ink wash style*. Click **✨ Suggest** to get one "
        "from your images."
    ),
    "object": (
        "Your **trigger** is a rare, unique token for the subject — e.g. *b3@rcup*, *sks dog*. "
        "Click **✨ Suggest** to get one from your images."
    ),
    "custom": (
        "Advanced: caption however you like. The **trigger** is whatever token or phrase you "
        "choose; **✨ Add captions** writes a literal description and appends your trigger, "
        "and blank captions fall back to it."
    ),
}
TRIGGER_PLACEHOLDER = {"style": "heavy impasto style", "object": "b3@rcup", "custom": "TOK"}

CAPTION_TIP = (
    "<details><summary>ℹ️ <b>How captioning works</b></summary>\n\n"
    "- For a <b>style</b>, captions describe only the <i>content</i> (subjects, layout, setting) "
    "and end with your style trigger — the model then learns the look, not the subjects.\n"
    "- For an <b>object / character</b>, captions describe the scene and tag the subject with your "
    "trigger token.\n"
    "- <b>✨ Auto-caption</b> fills these for you; edit anything you like. Blank captions fall "
    "back to the trigger.\n"
    "- Already have captions? Upload a <code>.txt</code> next to each image with the same name "
    "(e.g. <code>cat.jpg</code> + <code>cat.txt</code>) and they're filled in automatically.\n"
    "</details>"
)


def _paths(files) -> list[str]:
    return [p if isinstance(p, str) else getattr(p, "name", p) for p in (files or [])]


def _split_uploads(files):
    """Split a mixed upload into image paths and a {basename: caption} map from .txt sidecars.

    A .txt file whose name matches an image (e.g. cat.jpg + cat.txt) prefills that caption.
    """
    images, captions = [], {}
    for p in _paths(files):
        ext = os.path.splitext(p)[1].lower()
        base = os.path.splitext(os.path.basename(p))[0]
        if ext == ".txt":
            try:
                with open(p, encoding="utf-8") as f:
                    captions[base] = f.read().strip()
            except Exception:  # noqa: BLE001
                pass
        elif ext in IMAGE_EXTS:
            images.append(p)
    return images, captions


def _signin_state(profile: gr.OAuthProfile | None):
    # The form stays visible but locked (pointer-events:none) until sign-in: OAuth is a full-page
    # redirect, so anything typed before logging in would be wiped on the way back.
    if profile is None:
        return ("", gr.update(elem_classes=["locked"]), gr.update())
    return (
        "",
        gr.update(elem_classes=[]),
        gr.update(info=f"Becomes your model repo `{profile.username}/<name>` (+ a dataset repo). "
                       "Spaces and symbols are turned into dashes automatically."),
    )


def on_concept_change(concept: str):
    # picking a concept reveals the upload area (lora-ease "part by part" reveal)
    return (
        gr.update(visible=True),
        gr.update(placeholder=TRIGGER_PLACEHOLDER.get(concept, "")),
        TRIGGER_HELP.get(concept, ""),
    )


def load_captioning(files, trigger):
    """Reveal the captioning + training UI once images are uploaded; one row per image.

    If a matching `<name>.txt` sidecar was uploaded, its content prefills that image's caption.
    """
    paths, txt_caps = _split_uploads(files)
    n = len(paths)
    if n > MAX_IMAGES:
        raise gr.Error(f"For now, up to {MAX_IMAGES} images are supported (got {n}).")
    updates = [gr.update(visible=n > 0), gr.update(visible=n > 0)]  # captioning_area, post_upload
    for i in range(MAX_IMAGES):
        visible = i < n
        if visible:
            base = os.path.splitext(os.path.basename(paths[i]))[0]
            cap_val = txt_caps.get(base) or (trigger or "")
        else:
            cap_val = None
        updates.append(gr.update(visible=visible))  # row
        updates.append(gr.update(value=paths[i] if visible else None, visible=visible))  # image
        updates.append(gr.update(value=cap_val, visible=visible))  # caption
    return updates


def ai_suggest_trigger(files, concept_type):
    paths, _ = _split_uploads(files)
    if not paths:
        raise gr.Error("Upload images first.")
    return caption.suggest_trigger(paths, concept_type)


def ai_caption_all(files, concept_type, trigger):
    paths, _ = _split_uploads(files)
    if not paths:
        raise gr.Error("Upload images first.")
    trigger = (trigger or "").strip()
    if not trigger:
        trigger = caption.suggest_trigger(paths, concept_type)  # suggest one and reuse it everywhere
    caps = [caption.caption_one(p, concept_type, trigger) for p in paths]
    cap_updates = [gr.update(value=caps[i]) if i < len(caps) else gr.update() for i in range(MAX_IMAGES)]
    return [gr.update(value=trigger), *cap_updates]


def make_caption_one(idx):
    """Build a handler that captions only image `idx` (seeding the trigger if it's empty)."""
    def _caption_one(files, concept_type, trigger):
        paths, _ = _split_uploads(files)
        if idx >= len(paths):
            return gr.update(), gr.update()
        trigger = (trigger or "").strip()
        if not trigger:
            trigger = caption.suggest_trigger(paths, concept_type)
        return gr.update(value=trigger), gr.update(value=caption.caption_one(paths[idx], concept_type, trigger))
    return _caption_one


def gather_dataset(files, *captions):
    paths, _ = _split_uploads(files)
    return [[img, (captions[i] if i < len(captions) else "")] for i, img in enumerate(paths)]


def start_training(
    dataset_rows, concept_type, lora_name, trigger, rank, lora_alpha,
    max_train_steps, learning_rate, lr_scheduler, resolution, repeats, train_batch_size,
    gradient_accumulation_steps, seed, optimizer, use_8bit_adam, cache_latents,
    gradient_checkpointing, offload, quantization, lora_layers, make_gallery, num_gallery_images,
    custom_eval_prompts, flavor, timeout,
    profile: gr.OAuthProfile | None = None, oauth_token: gr.OAuthToken | None = None,
):
    if oauth_token is None or profile is None:
        return "❌ Please **sign in with Hugging Face** first (top-right).", "", ""
    if not dataset_rows:
        return "❌ Upload at least one image.", "", ""
    if not (lora_name or "").strip():
        return "❌ Give your LoRA a name.", "", ""

    image_paths = [r[0] for r in dataset_rows]
    captions = [r[1] for r in dataset_rows]
    params = {
        "concept_type": concept_type, "lora_name": lora_name,
        "instance_prompt": trigger,
        "rank": rank, "lora_alpha": lora_alpha, "max_train_steps": max_train_steps,
        "learning_rate": learning_rate, "lr_scheduler": lr_scheduler, "resolution": resolution,
        "repeats": repeats, "train_batch_size": train_batch_size,
        "gradient_accumulation_steps": gradient_accumulation_steps, "seed": seed,
        "optimizer": optimizer, "use_8bit_adam": bool(use_8bit_adam),
        "cache_latents": bool(cache_latents), "gradient_checkpointing": bool(gradient_checkpointing),
        "offload": bool(offload), "quantization": quantization, "lora_layers": lora_layers,
        "make_gallery": bool(make_gallery), "num_gallery_images": num_gallery_images,
        "custom_eval_prompts": custom_eval_prompts,
        "hf_token": oauth_token.token,
    }
    try:
        res = jobs.submit(params, image_paths, captions, 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 = (
        f"**Job:** [{res['job_id']}]({res['url']})  \n"
        f"**Dataset:** `{res['dataset_repo']}`  \n"
        f"**LoRA will be pushed to:** `{res['hub_model_id']}`  \n\n"
        f"Track progress in the **Monitor** tab (job id is prefilled)."
    )
    return status, link, res["job_id"]


def cost_estimate(steps, num_gallery, make_gallery, flavor):
    minutes, dollars = jobs.estimate_cost(steps, num_gallery, make_gallery, flavor)
    sps = jobs.SEC_PER_STEP.get(flavor, 2.35)
    return (
        f"💸 **Estimated cost:** ~${dollars:.2f} · ~{minutes:.0f} min on `{flavor}` "
        f"(~{sps:.1f}s/step, compiled). _Rough estimate — the first run also downloads the "
        f"gated Krea 2 weights (one-time ~2–3 min)._"
    )


def refresh(job_id, oauth_token: gr.OAuthToken | None = None):
    if not (job_id or "").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


with gr.Blocks(title="Krea 2 LoRA Trainer") as demo:
    gr.Markdown(
        "# 🎨 Krea 2 LoRA Trainer\n"
        "### Train a high-quality Krea 2 LoRA from your own images\n"
        "Trains on **Krea 2 RAW**, optionally validates on **Turbo**, runs on **HF Jobs**, "
        "pushed to your Hub. You only pay for the job's GPU minutes.",
        elem_id="main_title",
    )
    with gr.Row():
        gr.Column(scale=2, min_width=0)
        with gr.Column(scale=0, min_width=380):
            gr.LoginButton("Sign in with Hugging Face to use the UI", elem_id="login", elem_classes=["login_logout"])
            gr.Markdown("…or paste this to your coding agent:", elem_classes=["agent_hint"])
            gr.Code(f"curl {AGENT_URL}", language="shell", interactive=False,
                    show_label=False, elem_classes=["agent_curl"])
        gr.Column(scale=2, min_width=0)
    banner = gr.Markdown(elem_id="signin_banner")

    with gr.Column(elem_id="main_ui", elem_classes=["locked"]) as main_ui, gr.Tabs():
        with gr.Tab("Train"):
            lora_name = gr.Textbox(
                label="The name of your LoRA", placeholder="e.g. my-impasto-style",
                info="Has to be unique — used for your output model repo (you/<name>) and dataset repo.",
            )
            concept_type = gr.Radio(
                CONCEPT_CHOICES, value=None, label="What are you training?",
                info="Drives how images are captioned and what kind of trigger is suggested.",
            )

            with gr.Row(visible=False, equal_height=False) as image_upload:
                with gr.Column(scale=1):
                    images = gr.File(
                        label="Upload your images (4–30 ideal)", file_count="multiple",
                        file_types=["image", ".txt"], interactive=True, height=320,
                    )
                    gr.Markdown(
                        "_Already have captions? Drop a `.txt` next to each image with the same "
                        "name (e.g. `cat.png` + `cat.txt`) and they'll be filled in automatically. "
                        "Otherwise, you'll be able to caption your images here →_",
                        elem_id="upload_hint",
                    )
                with gr.Column(scale=3, visible=False, elem_id="captioning_area") as captioning_area:
                    gr.Markdown("## Trigger & captioning")
                    with gr.Row(elem_classes=["trigger_row"], equal_height=True):
                        trigger = gr.Textbox(
                            label="Trigger", placeholder=TRIGGER_PLACEHOLDER["style"],
                            scale=4, interactive=True,
                        )
                        suggest_btn = gr.Button("✨ Suggest", scale=1, min_width=120,
                                                variant="secondary", elem_classes=["suggest_btn"])
                    trigger_help = gr.Markdown(TRIGGER_HELP["style"])
                    autocaption_btn = gr.Button("✨ Add captions with Gemma",
                                                variant="primary")
                    gr.Markdown(CAPTION_TIP)
                    caption_rows, caption_imgs, caption_txts = [], [], []
                    for i in range(MAX_IMAGES):
                        with gr.Column(visible=False, elem_classes=["cap_item"]) as row:
                            img = gr.Image(
                                height=180, interactive=False, show_label=False,
                                container=False, elem_classes=["cap_img"],
                            )
                            with gr.Row(elem_classes=["cap_row"], equal_height=True):
                                cap = gr.Textbox(label=f"Caption {i + 1}", scale=14, interactive=True)
                                cap_btn = gr.Button("✨", scale=0, min_width=46,
                                                    variant="secondary", elem_classes=["cap_one_btn"])
                        cap_btn.click(make_caption_one(i), inputs=[images, concept_type, trigger],
                                      outputs=[trigger, cap])
                        caption_rows.append(row)
                        caption_imgs.append(img)
                        caption_txts.append(cap)

            with gr.Column(visible=False) as post_upload:
                with gr.Accordion("Advanced options", open=False, elem_classes=["accordion"]):
                    with gr.Row():
                        rank = gr.Number(label="LoRA rank", value=32, precision=0,
                                         info="Authors recommend 32; raise it for long runs / high-frequency styles.")
                        lora_alpha = gr.Number(label="LoRA alpha", value=32, precision=0,
                                               info="Keep equal to rank (scale 1.0).")
                    with gr.Row():
                        max_train_steps = gr.Number(label="Training steps", value=1000, precision=0)
                        learning_rate = gr.Number(label="Learning rate", value=3e-4,
                                                  info="3e-4 ~ 7e-4 with constant works well; go higher with cosine.")
                    with gr.Row():
                        lr_scheduler = gr.Dropdown(LR_SCHEDULERS, value="constant", label="LR scheduler")
                        resolution = gr.Number(label="Resolution", value=1024, precision=0)
                    with gr.Row():
                        repeats = gr.Number(label="Dataset repeats", value=1, precision=0)
                        seed = gr.Number(label="Seed", value=0, precision=0)
                    lora_layers = gr.Textbox(
                        label="Target layers (optional)", placeholder="wq,wk,wv,wo,gate",
                        info="Blank = the authors' recommended full set. For long runs, narrow to the "
                             "attention layers (wq,wk,wv,wo,gate) and raise the rank so prompt adherence holds.",
                    )
                    timeout = gr.Textbox(
                        label="Job timeout", value="6h",
                        info="Max job runtime before it's stopped. The first run also downloads the "
                             "gated Krea 2 weights, so leave headroom.",
                    )
                    with gr.Accordion("Memory / performance", open=False):
                        with gr.Row():
                            quantization = gr.Dropdown(QUANT_CHOICES, value="none", label="Quantization")
                            optimizer = gr.Dropdown(OPTIMIZERS, value="adamW", label="Optimizer")
                        with gr.Row():
                            use_8bit_adam = gr.Checkbox(label="8-bit Adam", value=True)
                            cache_latents = gr.Checkbox(label="Cache latents", value=True)
                        with gr.Row():
                            gradient_checkpointing = gr.Checkbox(label="Gradient checkpointing", value=True)
                            offload = gr.Checkbox(label="CPU offload (VAE + text encoder)", value=False)
                        with gr.Row():
                            train_batch_size = gr.Number(label="Batch size", value=1, precision=0)
                            gradient_accumulation_steps = gr.Number(label="Grad accumulation", value=1, precision=0)
                    with gr.Accordion("Preview gallery & README (on Turbo)", open=True,
                                      elem_classes=["accordion"]):
                        make_gallery = gr.Checkbox(
                            label="Generate a preview gallery + rich README", value=True,
                            info="After training, render sample images on Krea 2 Turbo with your LoRA "
                                 "and push a model-card README where each image is captioned by its prompt.",
                        )
                        num_gallery_images = gr.Number(label="Number of samples", value=3, precision=0)
                        custom_eval_prompts = gr.Textbox(
                            label="Custom showcase prompts (optional)", lines=4,
                            placeholder="One prompt per line — use <trigger> where the trigger should go.\n"
                                        "Leave blank to auto-generate diverse prompts with the LLM.",
                            info="If blank, the LLM writes diverse showcase prompts from your concept + trigger.",
                        )

                gr.Markdown("### Output & submit")
                flavor = gr.Dropdown(jobs.FLAVORS, value=jobs.DEFAULT_FLAVOR, label="GPU flavor")
                with gr.Accordion("GPU guide", open=False):
                    gr.Markdown(FLAVOR_GUIDE)

                cost_md = gr.Markdown(cost_estimate(1000, 3, True, jobs.DEFAULT_FLAVOR))
                submit_btn = gr.Button("🚀 Submit training job", variant="primary", size="lg")
                status = gr.Markdown()
                joblink = gr.Markdown()

        with gr.Tab("Monitor"):
            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=22, autoscroll=True, max_lines=22)

    dataset_state = gr.State([])

    caption_outputs = [captioning_area, post_upload]
    for r, im, c in zip(caption_rows, caption_imgs, caption_txts):
        caption_outputs += [r, im, c]

    demo.load(_signin_state, inputs=None, outputs=[banner, main_ui, lora_name])
    concept_type.change(
        on_concept_change, inputs=[concept_type], outputs=[image_upload, trigger, trigger_help],
    )
    images.change(load_captioning, inputs=[images, trigger], outputs=caption_outputs)
    suggest_btn.click(ai_suggest_trigger, inputs=[images, concept_type], outputs=[trigger])
    autocaption_btn.click(
        ai_caption_all, inputs=[images, concept_type, trigger], outputs=[trigger, *caption_txts],
    )

    submit_btn.click(
        gather_dataset, inputs=[images, *caption_txts], outputs=dataset_state,
    ).then(
        start_training,
        inputs=[dataset_state, concept_type, lora_name, trigger, rank, lora_alpha,
                max_train_steps, learning_rate, lr_scheduler, resolution, repeats, train_batch_size,
                gradient_accumulation_steps, seed, optimizer, use_8bit_adam, cache_latents,
                gradient_checkpointing, offload, quantization, lora_layers, make_gallery,
                num_gallery_images, custom_eval_prompts, flavor, timeout],
        outputs=[status, joblink, job_id],
    )
    cost_inputs = [max_train_steps, num_gallery_images, make_gallery, flavor]
    for comp in cost_inputs:
        comp.change(cost_estimate, inputs=cost_inputs, outputs=cost_md)

    refresh_btn.click(refresh, inputs=[job_id], outputs=[mon_status, mon_logs])


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