"""HF Jobs backend for the Krea 2 LoRA trainer Space.

Per training request the Space (cpu-basic, no GPU/torch):
  1. stages the uploaded images + a `metadata.jsonl` (per-image captions),
  2. pushes them to a private HF **dataset** repo under the signed-in user,
  3. generates a self-contained UV job script,
  4. submits it with `HfApi.run_uv_job(... token=<user>)` so the job runs + is billed
     to the signed-in user and the trained LoRA is pushed to their Hub.

Token split (important):
  * The **gated Krea 2 weights** (`krea/Krea-2-Raw`, `krea/Krea-2-Turbo`) are not public.
    They are pre-downloaded *inside the job* with the Space's `KREA_TOKEN` secret and passed
    to the trainer as **local dirs**, so `from_pretrained` needs no Krea auth.
  * Everything else (dataset download, `create_repo`/`upload_folder` of the LoRA) uses the
    job's ambient `HF_TOKEN` env = the **signed-in user's** token. The Krea token never
    touches the user's repos and the user's token never needs Krea access.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import tempfile
from pathlib import Path

from huggingface_hub import HfApi

# The Krea 2 LoRA trainer (diffusers PR #14046) is merged to `main`. Override with the
# `DIFFUSERS_REF` Space variable to pin a release tag/commit if needed.
DIFFUSERS_REF = os.environ.get("DIFFUSERS_REF", "main")
BASE_MODEL_RAW = "krea/Krea-2-Raw"      # non-distilled base — train LoRA on this
BASE_MODEL_TURBO = "krea/Krea-2-Turbo"  # 8-step distilled — validate / infer on this

DEFAULT_FLAVOR = "rtx-pro-6000"
FLAVORS = [
    ("RTX PRO 6000 — 96 GB ($0.046/min)", "rtx-pro-6000"),
    ("A100 — 80 GB ($0.042/min)", "a100-large"),
]
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}

# Cost model. Prices are per-minute of actual runtime (see FLAVORS labels). Timing is empirical on
# RTX PRO 6000 with regional compile (compile_repeated_blocks): ~2.35 s/step steady-state at rank 32
# / res 1024; A100 is a bit slower. FIXED_SEC = scheduling + 62 GB gated-weight download + dataset +
# showcase model load. Each Turbo gallery image (8 steps) ~6 s.
PRICE_PER_MIN = {"rtx-pro-6000": 0.046, "a100-large": 0.042}
SEC_PER_STEP = {"rtx-pro-6000": 2.35, "a100-large": 2.7}
FIXED_SEC = 240
SEC_PER_GALLERY_IMG = 6


def estimate_cost(steps, num_gallery, make_gallery, flavor):
    """Return (minutes, dollars) for a run. Rough — first run also pays a one-time weight download."""
    steps = max(0, int(steps or 0))
    ng = max(0, int(num_gallery or 0)) if make_gallery else 0
    sps = SEC_PER_STEP.get(flavor, 2.35)
    price = PRICE_PER_MIN.get(flavor, 0.046)
    total_sec = FIXED_SEC + steps * sps + ng * SEC_PER_GALLERY_IMG
    minutes = total_sec / 60.0
    return minutes, minutes * price

# Deterministic on-job paths the trainer reads from (baked into the CLI args below).
JOB_RAW = "/tmp/krea/raw"
JOB_TURBO = "/tmp/krea/turbo"
JOB_DATA = "/tmp/data"
JOB_OUT = "/tmp/out"
JOB_BNB = "/tmp/bnb.json"


def slug(name: str) -> str:
    s = re.sub(r"[^a-zA-Z0-9-]+", "-", (name or "").strip()).strip("-").lower()
    return s or "krea2-lora"


def _namespace(token: str) -> str:
    from huggingface_hub import whoami  # noqa: PLC0415
    return whoami(token=token)["name"]


def _custom_prompts(params: dict, instance_prompt: str) -> list[str]:
    """Parse the optional custom eval-prompts textarea (one per line). `<trigger>` placeholders are
    substituted with the actual trigger. Empty → []  → the job asks the LLM (or falls back)."""
    raw = (params.get("custom_eval_prompts") or "").strip()
    if not raw:
        return []
    trig = (instance_prompt or "").strip()
    out = []
    for line in raw.splitlines():
        line = line.strip()
        if not line:
            continue
        for ph in ("<trigger phrase>", "<trigger>", "<TOK>", "TOK"):
            line = line.replace(ph, trig)
        out.append(line)
    return out


def build_metadata(image_paths: list[str], captions: list[str], instance_prompt: str) -> list[dict]:
    """One `metadata.jsonl` row per image: {file_name, prompt}. Empty captions fall back to the
    instance prompt (the trigger sentence). Files are renamed to a stable `0000.ext` order."""
    rows = []
    fallback = (instance_prompt or "a photo").strip()
    for i, p in enumerate(image_paths):
        cap = ""
        if i < len(captions) and captions[i]:
            cap = str(captions[i]).strip()
        rows.append({"file_name": f"{i:04d}{Path(p).suffix.lower()}", "prompt": cap or fallback})
    return rows


def build_train_args(params: dict, hub_model_id: str) -> list[str]:
    """Turn UI params into the `train_dreambooth_lora_krea2.py` CLI. Krea weights are passed as
    local dirs (pre-downloaded in the job); the dataset is a local imagefolder (image/prompt cols)."""
    instance_prompt = (params.get("instance_prompt") or "TOK").strip()
    args = [
        "--pretrained_model_name_or_path", JOB_RAW,
        "--validation_model_path", JOB_TURBO,
        "--dataset_name", JOB_DATA,
        "--image_column", "image",
        "--caption_column", "prompt",
        "--instance_prompt", instance_prompt,
        "--output_dir", JOB_OUT,
        "--mixed_precision", "bf16",
        "--resolution", str(int(params["resolution"])),
        "--train_batch_size", str(int(params["train_batch_size"])),
        "--gradient_accumulation_steps", str(int(params["gradient_accumulation_steps"])),
        "--repeats", str(int(params["repeats"])),
        "--rank", str(int(params["rank"])),
        "--lora_alpha", str(int(params["lora_alpha"])),
        "--learning_rate", str(float(params["learning_rate"])),
        "--lr_scheduler", str(params["lr_scheduler"]),
        "--lr_warmup_steps", "0",
        "--max_train_steps", str(int(params["max_train_steps"])),
        "--optimizer", str(params["optimizer"]),
        "--seed", str(int(params["seed"])),
        "--push_to_hub",
        "--hub_model_id", hub_model_id,
    ]
    # Previews are produced by our own post-training showcase step (a multi-prompt gallery + rich
    # README), so the trainer itself always skips final inference.
    args += ["--skip_final_inference"]
    if params.get("lora_layers"):
        args += ["--lora_layers", str(params["lora_layers"]).strip()]
    if params.get("gradient_checkpointing", True):
        args += ["--gradient_checkpointing"]
    if params.get("cache_latents", True):
        args += ["--cache_latents"]
    if params.get("offload"):
        args += ["--offload"]
    if params.get("use_8bit_adam") and str(params["optimizer"]).lower() == "adamw":
        args += ["--use_8bit_adam"]
    quant = params.get("quantization", "none")
    if quant == "fp8":
        args += ["--do_fp8_training"]
    elif quant == "4bit":
        args += ["--bnb_quantization_config_path", JOB_BNB]
    return args


# --------------------------------------------------------------------------------------
# UV job script (runs on HF Jobs GPU hardware)
# --------------------------------------------------------------------------------------

# The script is rendered in two pieces: a small `.format()`-ed HEADER (PEP 723 metadata + config
# constants) and a verbatim BODY (raw string, never `.format()`-ed) so the showcase code's many
# dict / set literals don't need brace-escaping.
JOB_HEADER = '''# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "diffusers @ git+https://github.com/huggingface/diffusers.git@{ref}",
#   "torch",
#   "torchvision",
#   "transformers>=4.41.2",
#   "accelerate>=0.31.0",
#   "peft>=0.11.1",
#   "datasets",
#   "bitsandbytes",
#   "prodigyopt",
#   "ftfy",
#   "sentencepiece",
#   "hf_transfer",
#   "huggingface_hub[hf-xet]",
# ]
# ///
"""Auto-generated Krea 2 DreamBooth-LoRA job: trains on RAW, then renders a Turbo preview gallery."""
REF = "{ref}"
RAW, TURBO, DATA, BNB, OUT = "{raw}", "{turbo}", "{data}", "{bnb}", "{out}"
DATASET_REPO = {dataset_repo!r}
QUANT = {quant!r}
CONCEPT = {concept!r}
TRIGGER = {trigger!r}
MAKE_GALLERY = {make_gallery}
NUM_GALLERY = {num_gallery}
CUSTOM_PROMPTS = {custom_prompts}
HUB_MODEL_ID = {hub_model_id!r}
TRAIN_ARGS = {train_args}
'''

JOB_BODY = r'''
import json, os, re, subprocess, sys, urllib.request
from pathlib import Path

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
KREA_TOKEN = os.environ.get("KREA_TOKEN") or None        # Krea weights (public now; token optional)
CAPTION_TOKEN = os.environ.get("CAPTION_HF_TOKEN", "")   # LLM eval-prompt generation (optional)

BASE_RAW_ID, TURBO_ID = "krea/Krea-2-Raw", "krea/Krea-2-Turbo"
CAPTION_MODEL, CAPTION_PROVIDER = "google/gemma-4-31B-it", "novita"
SCRIPT_URL = ("https://raw.githubusercontent.com/huggingface/diffusers/"
              + REF + "/examples/dreambooth/train_dreambooth_lora_krea2.py")


def gen_eval_prompts(n):
    """Ask the LLM for n diverse showcase prompts that include the trigger."""
    from huggingface_hub import InferenceClient
    if CONCEPT == "style":
        instr = (
            "Generate " + str(n) + " diverse, vivid text-to-image prompts to showcase an image "
            "style. Each prompt must describe a COMPLETELY different subject and scene (an animal "
            "in a landscape, a vehicle on a road, a person doing an action, an interior, etc.), be "
            "1-2 sentences, and MUST end with this exact phrase: " + repr(TRIGGER) + ". Return ONLY "
            "a JSON array of " + str(n) + " strings."
        )
    else:
        instr = (
            "Generate " + str(n) + " diverse, vivid text-to-image prompts to showcase a specific "
            "subject. Each prompt must place the subject in a COMPLETELY different scene/context, "
            "be 1-2 sentences, and MUST include this exact token: " + repr(TRIGGER) + ". Return ONLY "
            "a JSON array of " + str(n) + " strings."
        )
    client = InferenceClient(model=CAPTION_MODEL, provider=CAPTION_PROVIDER, token=CAPTION_TOKEN)
    r = client.chat_completion(messages=[{"role": "user", "content": instr}],
                               max_tokens=600, temperature=0.8)
    txt = (r.choices[0].message.content or "").strip()
    m = re.search(r"\[.*\]", txt, re.S)
    if m:
        prompts = [str(p).strip() for p in json.loads(m.group(0)) if str(p).strip()]
    else:
        prompts = [ln.strip().lstrip("-*0123456789. ").strip(' "') for ln in txt.splitlines()]
        prompts = [p for p in prompts if p]
    return prompts[:n] or [TRIGGER]


def fallback_prompts():
    if CONCEPT == "style":
        return [
            "A golden labrador running across a sunny farm, a pickup truck on a dusty road far behind, " + TRIGGER,
            "A fishing boat moored in a narrow canal between tall old buildings, " + TRIGGER,
            "A deer grazing in a dense forest with the bright sun overhead, " + TRIGGER,
        ]
    return [
        "A photo of " + TRIGGER + " on a wooden table indoors.",
        "A photo of " + TRIGGER + " outdoors on a patch of grass.",
        "A close-up photo of " + TRIGGER + " against a plain background.",
    ]


def esc(s):
    return str(s).replace("\\", "\\\\").replace('"', '\\"')


def write_readme(rows):
    widget = "\n".join('- text: "' + esc(p) + '"\n  output:\n    url: ' + fn for fn, p in rows)
    gallery = "\n".join("![sample](./" + fn + ')\n\n> *"' + p.replace('"', "'") + '"*\n' for fn, p in rows)
    trig = ("the phrase `" + TRIGGER + "`") if CONCEPT == "style" else ("the token `" + TRIGGER + "`")
    md = (
        "---\n"
        "base_model: " + BASE_RAW_ID + "\n"
        "tags:\n- text-to-image\n- diffusers\n- lora\n- krea2\n- template:sd-lora\n"
        "license: apache-2.0\n"
        "instance_prompt: " + json.dumps(TRIGGER) + "\n"
        "widget:\n" + widget + "\n"
        "---\n\n"
        "# Krea 2 LoRA — " + HUB_MODEL_ID + "\n\n"
        "<Gallery />\n\n"
        "A DreamBooth-LoRA for **Krea 2**, trained on **Krea 2 RAW** and shown on **Krea 2 Turbo**. "
        "The samples below were generated with this LoRA on Turbo (8 steps).\n\n"
        "## Trigger\n\nUse " + trig + " to invoke the concept.\n\n"
        "## Samples\n\n" + gallery + "\n"
        "## Use it with diffusers\n\n"
        "```py\n"
        "import torch\n"
        "from diffusers import Krea2Pipeline\n\n"
        'pipe = Krea2Pipeline.from_pretrained("' + TURBO_ID + '", torch_dtype=torch.bfloat16).to("cuda")\n'
        'pipe.load_lora_weights("' + HUB_MODEL_ID + '")\n'
        'image = pipe("' + esc(rows[0][1]) + '", num_inference_steps=8, guidance_scale=0.0).images[0]\n'
        'image.save("output.png")\n'
        "```\n"
    )
    Path(os.path.join(OUT, "README.md")).write_text(md)


def showcase():
    import torch
    from diffusers import Krea2Pipeline
    from huggingface_hub import HfApi

    prompts = list(CUSTOM_PROMPTS or [])
    if not prompts and CAPTION_TOKEN:
        try:
            prompts = gen_eval_prompts(NUM_GALLERY)
        except Exception as e:
            print("=== eval-prompt generation failed, using fallbacks ===", repr(e), flush=True)
    if not prompts:
        prompts = fallback_prompts()
    prompts = prompts[:max(1, NUM_GALLERY)]

    print("=== showcase: load Turbo + trained LoRA, render " + str(len(prompts)) + " samples ===", flush=True)
    pipe = Krea2Pipeline.from_pretrained(TURBO, torch_dtype=torch.bfloat16).to("cuda")
    pipe.load_lora_weights(OUT)

    rows = []
    for i, p in enumerate(prompts):
        print(">>> sample " + str(i) + ": " + p, flush=True)
        image = pipe(p, num_inference_steps=8, guidance_scale=0.0).images[0]
        fn = "sample_" + str(i) + ".png"
        image.save(os.path.join(OUT, fn))
        rows.append((fn, p))

    write_readme(rows)

    api = HfApi()  # ambient HF_TOKEN = user
    for fn, _ in rows:
        api.upload_file(path_or_fileobj=os.path.join(OUT, fn), path_in_repo=fn, repo_id=HUB_MODEL_ID)
    api.upload_file(path_or_fileobj=os.path.join(OUT, "README.md"), path_in_repo="README.md", repo_id=HUB_MODEL_ID)
    print("=== showcase: pushed gallery + README to " + HUB_MODEL_ID + " ===", flush=True)


def patch_trainer_compile(path):
    """Inject regional torch.compile (compile_repeated_blocks) into the fetched trainer for a ~1.84x
    LoRA training speedup on Krea2 (4.27 -> 2.32 s/it, ~12s warmup). Guarded at runtime so any
    failure falls back to uncompiled training."""
    src = Path(path).read_text()
    needle = "transformer.add_adapter(transformer_lora_config)"
    if needle not in src:
        print("=== compile patch skipped: anchor not found ===", flush=True)
        return
    inject = (
        needle + "\n"
        "    try:\n"
        "        import torch._dynamo as _dyn\n"
        "        _dyn.config.cache_size_limit = 64\n"
        "        transformer.compile_repeated_blocks(fullgraph=True)\n"
        "        print('=== compile_repeated_blocks ENABLED ===', flush=True)\n"
        "    except Exception as _e:\n"
        "        print('compile_repeated_blocks skipped:', _e, flush=True)\n"
    )
    Path(path).write_text(src.replace(needle, inject, 1))
    print("=== trainer patched: compile_repeated_blocks ===", flush=True)


def main():
    from huggingface_hub import snapshot_download

    print("=== 1/4 download Krea 2 weights ===", flush=True)
    snapshot_download(BASE_RAW_ID, local_dir=RAW, token=KREA_TOKEN)
    snapshot_download(TURBO_ID, local_dir=TURBO, token=KREA_TOKEN)

    print("=== 2/4 download dataset (user token / HF_TOKEN env) ===", flush=True)
    snapshot_download(DATASET_REPO, repo_type="dataset", local_dir=DATA)

    if QUANT == "4bit":
        Path(BNB).write_text(json.dumps({
            "load_in_4bit": True, "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_compute_dtype": "bfloat16",
        }))

    print("=== 3/4 fetch trainer script @ " + REF + " ===", flush=True)
    urllib.request.urlretrieve(SCRIPT_URL, "/tmp/train_dreambooth_lora_krea2.py")
    patch_trainer_compile("/tmp/train_dreambooth_lora_krea2.py")

    print("=== 4/4 accelerate launch (pushes LoRA to the Hub) ===", flush=True)
    cmd = [sys.executable, "-m", "accelerate.commands.launch",
           "/tmp/train_dreambooth_lora_krea2.py", *TRAIN_ARGS]
    print(">>> " + " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)

    if MAKE_GALLERY:
        try:
            showcase()
        except Exception as e:
            print("=== showcase failed (LoRA already trained + pushed) ===", repr(e), flush=True)

    print("=== DONE ===", flush=True)


if __name__ == "__main__":
    main()
'''


# --------------------------------------------------------------------------------------
# Standalone, CLI-driven version of the same job — handed to external agents so they can
# train on their OWN HF account with only their own token (Krea 2 weights are public).
# It sets the exact module-level constants JOB_BODY expects, but from argparse instead of
# baked .format() literals, then reuses JOB_BODY verbatim (single source of truth).
# --------------------------------------------------------------------------------------
JOB_AGENT_PEP723 = '''# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "diffusers @ git+https://github.com/huggingface/diffusers.git@{ref}",
#   "torch",
#   "torchvision",
#   "transformers>=4.41.2",
#   "accelerate>=0.31.0",
#   "peft>=0.11.1",
#   "datasets",
#   "bitsandbytes",
#   "prodigyopt",
#   "ftfy",
#   "sentencepiece",
#   "hf_transfer",
#   "huggingface_hub[hf-xet]",
# ]
# ///
"""Standalone Krea 2 DreamBooth-LoRA job (CLI-driven). Run with:
    hf jobs uv run --flavor rtx-pro-6000 --timeout 1h -s HF_TOKEN=$HF_TOKEN \\
        train_job.py -- --dataset <repo> --lora-name <name> --trigger "<trigger>"
"""
REF = "{ref}"
'''

JOB_AGENT_ARGPARSE = r'''
import argparse
import os as _os

RAW, TURBO, DATA, BNB, OUT = "/tmp/krea/raw", "/tmp/krea/turbo", "/tmp/data", "/tmp/bnb.json", "/tmp/out"

_p = argparse.ArgumentParser(description="Train a Krea 2 DreamBooth-LoRA on HF Jobs.")
_p.add_argument("--dataset", required=True, help="HF dataset repo with 'image' + 'prompt' columns.")
_p.add_argument("--lora-name", required=True, help="Output name; pushed to <you>/<name>.")
_p.add_argument("--trigger", required=True, help="Trigger phrase (style) or token (subject/face).")
_p.add_argument("--concept", default="style", choices=["style", "character", "object", "face"])
_p.add_argument("--hub-model-id", default="", help="Full repo id; default <you>/<lora-name>.")
_p.add_argument("--steps", type=int, default=1000)
_p.add_argument("--rank", type=int, default=32)
_p.add_argument("--lora-alpha", type=int, default=32)
_p.add_argument("--learning-rate", type=float, default=3e-4)
_p.add_argument("--lr-scheduler", default="constant")
_p.add_argument("--resolution", type=int, default=1024)
_p.add_argument("--repeats", type=int, default=1)
_p.add_argument("--train-batch-size", type=int, default=1)
_p.add_argument("--gradient-accumulation-steps", type=int, default=1)
_p.add_argument("--optimizer", default="adamW")
_p.add_argument("--seed", type=int, default=0)
_p.add_argument("--quantization", default="none", choices=["none", "fp8", "4bit"])
_p.add_argument("--no-gradient-checkpointing", action="store_true")
_p.add_argument("--no-cache-latents", action="store_true")
_p.add_argument("--use-8bit-adam", action="store_true")
_p.add_argument("--no-gallery", action="store_true", help="Skip the preview gallery + rich README.")
_p.add_argument("--num-gallery", type=int, default=3)
_a = _p.parse_args()

DATASET_REPO = _a.dataset
CONCEPT = _a.concept
TRIGGER = _a.trigger.strip()
QUANT = _a.quantization
MAKE_GALLERY = not _a.no_gallery
NUM_GALLERY = _a.num_gallery
CUSTOM_PROMPTS = []

from huggingface_hub import whoami as _whoami
_ns = _whoami(token=_os.environ["HF_TOKEN"])["name"]
HUB_MODEL_ID = _a.hub_model_id.strip() or (_ns + "/" + _a.lora_name)

TRAIN_ARGS = [
    "--pretrained_model_name_or_path", RAW,
    "--validation_model_path", TURBO,
    "--dataset_name", DATA,
    "--image_column", "image",
    "--caption_column", "prompt",
    "--instance_prompt", TRIGGER,
    "--output_dir", OUT,
    "--mixed_precision", "bf16",
    "--resolution", str(_a.resolution),
    "--train_batch_size", str(_a.train_batch_size),
    "--gradient_accumulation_steps", str(_a.gradient_accumulation_steps),
    "--repeats", str(_a.repeats),
    "--rank", str(_a.rank),
    "--lora_alpha", str(_a.lora_alpha),
    "--learning_rate", str(_a.learning_rate),
    "--lr_scheduler", _a.lr_scheduler,
    "--lr_warmup_steps", "0",
    "--max_train_steps", str(_a.steps),
    "--optimizer", _a.optimizer,
    "--seed", str(_a.seed),
    "--push_to_hub",
    "--hub_model_id", HUB_MODEL_ID,
    "--skip_final_inference",
]
if not _a.no_gradient_checkpointing:
    TRAIN_ARGS += ["--gradient_checkpointing"]
if not _a.no_cache_latents:
    TRAIN_ARGS += ["--cache_latents"]
if _a.use_8bit_adam and _a.optimizer.lower() == "adamw":
    TRAIN_ARGS += ["--use_8bit_adam"]
if QUANT == "fp8":
    TRAIN_ARGS += ["--do_fp8_training"]
elif QUANT == "4bit":
    TRAIN_ARGS += ["--bnb_quantization_config_path", BNB]
'''


def build_agent_job_script() -> str:
    """The full standalone UV script committed as train_job.py for agents to curl + submit."""
    return JOB_AGENT_PEP723.format(ref=DIFFUSERS_REF) + JOB_AGENT_ARGPARSE + JOB_BODY


def submit(params: dict, image_paths: list[str], captions: list[str],
           flavor: str, timeout: str) -> dict:
    """Stage dataset → push private dataset repo → generate UV script → submit job.
    Returns {job_id, url, dataset_repo, hub_model_id}."""
    token = (params.get("hf_token") or "").strip()
    if not token:
        raise ValueError("Missing user token (sign in with Hugging Face).")
    if not os.environ.get("KREA_TOKEN"):
        raise RuntimeError("Space is missing the KREA_TOKEN secret (gated Krea 2 access).")

    ns = _namespace(token)
    name = slug(params.get("lora_name", ""))
    dataset_repo = f"{ns}/{name}-dataset"
    hub_model_id = (params.get("hub_model_id") or "").strip() or f"{ns}/{name}"

    api = HfApi(token=token)
    tmp = Path(tempfile.mkdtemp(prefix="krea2-"))
    try:
        # stage images under stable names + metadata.jsonl
        rows = build_metadata(image_paths, captions, params.get("instance_prompt", ""))
        for row, src in zip(rows, image_paths):
            shutil.copy(src, tmp / row["file_name"])
        (tmp / "metadata.jsonl").write_text(
            "\n".join(json.dumps(r) for r in rows) + "\n"
        )

        # push the dataset (private, user namespace)
        api.create_repo(dataset_repo, repo_type="dataset", private=True, exist_ok=True, token=token)
        api.upload_folder(repo_id=dataset_repo, repo_type="dataset", folder_path=str(tmp), token=token)

        # render the job script (formatted header + verbatim body)
        train_args = build_train_args(params, hub_model_id)
        custom_prompts = _custom_prompts(params, instance_prompt=params.get("instance_prompt", ""))
        header = JOB_HEADER.format(
            ref=DIFFUSERS_REF, raw=JOB_RAW, turbo=JOB_TURBO, data=JOB_DATA, bnb=JOB_BNB, out=JOB_OUT,
            dataset_repo=dataset_repo, quant=params.get("quantization", "none"),
            concept=params.get("concept_type", "style"),
            trigger=(params.get("instance_prompt") or "").strip(),
            make_gallery=bool(params.get("make_gallery", True)),
            num_gallery=int(params.get("num_gallery_images", 3) or 3),
            custom_prompts=repr(custom_prompts),
            hub_model_id=hub_model_id, train_args=json.dumps(train_args),
        )
        script_path = tmp / "job_train.py"
        script_path.write_text(header + JOB_BODY)

        # HF_TOKEN = user (push + dataset); KREA_TOKEN = gated Krea weights only;
        # CAPTION_HF_TOKEN = LLM eval-prompt generation for the showcase (optional).
        secrets = {"HF_TOKEN": token, "KREA_TOKEN": os.environ["KREA_TOKEN"]}
        if os.environ.get("CAPTION_HF_TOKEN"):
            secrets["CAPTION_HF_TOKEN"] = os.environ["CAPTION_HF_TOKEN"]

        job = api.run_uv_job(
            str(script_path),
            flavor=flavor,
            timeout=timeout,
            secrets=secrets,
            token=token,
        )
        job_id = getattr(job, "id", "") or ""
        url = getattr(job, "url", "") or (f"https://huggingface.co/jobs/{ns}/{job_id}" if job_id else "")
        return {"job_id": job_id, "url": url, "dataset_repo": dataset_repo, "hub_model_id": hub_model_id}
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def job_logs(job_id: str, token: str = "") -> str:
    try:
        return "\n".join(HfApi(token=token).fetch_job_logs(job_id=job_id, token=token))
    except Exception as e:  # noqa: BLE001
        return f"(could not fetch logs: {e})"


def job_status(job_id: str, token: str = "") -> str:
    try:
        job = HfApi(token=token).inspect_job(job_id=job_id, token=token)
        status = getattr(job, "status", None)
        stage = getattr(status, "stage", None)
        if stage is None and isinstance(status, dict):
            stage = status.get("stage")
        return str(stage or status or "UNKNOWN")
    except Exception as e:  # noqa: BLE001
        return f"UNKNOWN ({e})"