server.py

"""Workshop - LoRA Pipeline Configurator.



Configure a Wan 2.2 I2V base pipeline, attach LoRAs from the reference

library (warehouse repo), label example image+prompt pairs, then publish a

self-contained private Space that pulls its LoRA from the warehouse on first

generate (cached afterwards). A conversational assistant (HF Inference, using

the Space's HF_TOKEN secret) drives the skills by emitting small JSON tool

calls that this server parses and executes.

"""
import json
import os
import re
import shutil
import tempfile
import time
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from huggingface_hub import HfApi, SpaceHardware, upload_folder, hf_hub_download
import requests

# ---------------------------------------------------------------------------
# Space configuration
# ---------------------------------------------------------------------------
OWNER = "00000tt"
WAREHOUSE_REPO = os.environ.get("WAREHOUSE_REPO", "00000tt/wan22-loras")
SANDBOX = Path(os.environ.get("SANDBOX_DIR", "/tmp/workshop_build"))
EXAMPLES_DIR = SANDBOX / "examples"

# Base pipeline recipes the configurator can offer. Proven on the owner's
# example Space (Wan 2.2 I2V Lightning 14B, FP8 + AoT, 6 steps).
BASE_MODELS = {
    "Wan 2.2 I2V Lightning 14B (AoT FP8)": {
        "pipeline_id": "WanImageToVideoPipeline",
        "repo_id": "TestOrganizationPleaseIgnore/WAMU_v2_WAN2.2_I2V_LIGHTNING",
        "steps": 6,
        "guidance": 1.0,
        "quant": "fp8",
        "max_frames": 241,
        "template": "wan22_i2v",
    },
    "Qwen-Image-Edit-2511 LoRAs Fast": {
        "pipeline_id": "QwenImageEditPlusPipeline",
        "repo_id": "Qwen/Qwen-Image-Edit-2511",
        "steps": 30,
        "guidance": 5.0,
        "quant": "bf16",
        "max_frames": 0,
        "template": "qwen_image_edit_2511",
        "multi_image": True,
    },
    "Wan 2.2 I2V Lightning Studio (Example-driven)": {
        "pipeline_id": "WanImageToVideoPipeline",
        "repo_id": "TestOrganizationPleaseIgnore/WAMU_v2_WAN2.2_I2V_LIGHTNING",
        "steps": 6,
        "guidance": 1.0,
        "quant": "fp8",
        "max_frames": 241,
        "template": "wan22_i2v_studio",
        "multi_image": True,
    },
    "Qwen Image Edit Object Manipulator (2-photo)": {
        "pipeline_id": "QwenImageEditPlusPipeline",
        "repo_id": "Qwen/Qwen-Image-Edit-2509",
        "steps": 4,
        "guidance": 1.0,
        "quant": "bf16",
        "max_frames": 0,
        "template": "qwen_object_manip",
        "multi_image": True,
    },
}

CHAT_MODELS = [
    "Qwen/Qwen2.5-0.5B-Instruct",
    "google/gemma-2-2b-it",
    "HuggingFaceH4/zephyr-7b-beta",
]
CHAT_TIMEOUT = 60

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_headers=["*"],
    allow_methods=["*"],
)


class State:
    """Holds the vault token and the live pipeline config."""

    vault_token: Optional[str] = os.environ.get("HF_TOKEN") or None
    config: dict = {}


def _token() -> str:
    return State.vault_token


def _safe_whoami() -> Optional[str]:
    if not State.vault_token:
        return None
    try:
        return HfApi(token=State.vault_token).whoami().get("name")
    except Exception:
        return None


def _slugify(name: str) -> str:
    slug = re.sub(r"[^a-z0-9\-_]", "-", name.lower()).strip("-")
    return slug or "wan-lora-studio"


# ---------------------------------------------------------------------------
# Warehouse / reference library
# ---------------------------------------------------------------------------
@app.get("/api/warehouse")
def warehouse_list():
    """Live list of loose .safetensors in the reference LoRA library."""
    if not State.vault_token:
        return {"status": "vault-empty", "repo": WAREHOUSE_REPO, "assets": []}
    try:
        api = HfApi(token=State.vault_token)
        sizes: dict[str, int] = {}
        try:
            for f in api.list_repo_tree(WAREHOUSE_REPO, repo_type="model", recursive=True):
                name = getattr(f, "path", None) or getattr(f, "rfilename", "")
                if name.endswith(".safetensors"):
                    sizes[name] = int(getattr(f, "size", 0) or 0)
        except Exception:
            sizes = {}
        files = api.list_repo_files(WAREHOUSE_REPO, repo_type="model")
        assets = [
            {"name": f, "size": int(sizes.get(f) or 0)}
            for f in files if f.endswith(".safetensors")
        ]
        return {"assets": sorted(assets, key=lambda a: a["name"]), "repo": WAREHOUSE_REPO}
    except Exception as e:
        return {"assets": [], "repo": WAREHOUSE_REPO, "error": str(e)}


# ---------------------------------------------------------------------------
# One-button "bowl": photos + prompt + LoRA -> get examples, green tick,
# push one button, whole packet installed as a private Space. Done.
# ---------------------------------------------------------------------------
LORA_HINTS = {
    # default motion prompt per LoRA family (editable in the bowl, then kept)
    "WAN-2.2-I2V-FaceDownAssUp-HIGH-v1.safetensors": (
        "Face down, ass up on the bed, she pushes her hips up high, "
        "pulls her head back and looks over her shoulder at the camera, "
        "slow zoom in, warm cinematic light, steady hands, natural motion, skin detail"
    ),
    "WAN-2.2-I2V-FaceDownAssUp-LOW-v1.safetensors": (
        "Face down on the bed, low hips, she lifts her head, arching slowly, "
        "calm movement, sidelight, gentle breathing motion, cinematic soft light"
    ),
    "WAN-2.2-I2V-POV-Cowgirl-HIGH-v1.0-fixed.safetensors": (
        "First-person view, girl on top riding hard, her hips roll fast, "
        "hands on the camera, she tilts her head back, breathy motion, "
        "high angle energy, intimate cinematic tones"
    ),
    "WAN-2.2-I2V-POV-Cowgirl-LOW-v1.0-fixed.safetensors": (
        "First-person view, girl on top, slow grind, her body rocks gently, "
        "she leans in close, eyes on camera, soft motion, warm intimate light"
    ),
}
DEFAULT_HINT = (
    "camera holds on the subject, gentle natural motion, she looks at the "
    "camera and smiles, slow movement, cinematic warm light, high detail"
)


def _hint_for_lora(name: str) -> str:
    if not name:
        return DEFAULT_HINT
    for key, hint in LORA_HINTS.items():
        if key in name:
            return hint
    return DEFAULT_HINT


@app.get("/api/hints")
def hints_list():
    """Per-LoRA example prompts the workshop can pre-fill."""
    return {"hints": {k: v[:160] for k, v in LORA_HINTS.items()},
            "default": DEFAULT_HINT,
            "adapters": [
                "Qwen-Image-Edit-2511-Object-Adder",
                "Qwen-Image-Edit-2511-Object-Remover",
                "QIE-2511-Object-Remover-v2",
                "Zoom-Master",
                "Extract-Outfit",
                "Outfit-Design-Layout",
            ]}


@app.post("/api/bowl")
async def bowl_send(files: list[UploadFile] = File(...),

                    prompt: str = Form(""),

                    lora: str = Form(""),

                    base: str = Form("")):
    """The one-button flow: photos + optional prompt + optional LoRA

    -> label one example AND publish a whole private Space, no further steps.

    """
    token = _token()
    if not token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    if not files:
        raise HTTPException(status_code=400, detail="Drop at least one photo into the bowl.")
    prompt = (prompt or "").strip() or _hint_for_lora(lora)
    if not lora:
        assets = (warehouse_list().get("assets") or [])
        lora = assets[0]["name"] if assets else ""
    lora = lora or ""

    # label example
    n_ex = len(list(Path(EXAMPLES_DIR).glob("example_*"))) + 1
    ex_dir = Path(EXAMPLES_DIR) / f"example_{n_ex}"
    ex_dir.mkdir(parents=True, exist_ok=True)
    rel = []
    for i, f in enumerate(files, start=1):
        data = await f.read()
        ext = (f.filename or "").rsplit(".", 1)[-1].lower()
        ext = ext if ext in ("jpg", "jpeg", "png", "webp") else "jpg"
        rel_name = f"image_{i}.{ext}"
        (ex_dir / rel_name).write_bytes(data)
        rel.append(rel_name)
    (ex_dir / "prompt.txt").write_text(prompt[:220], encoding="utf-8")
    if lora:
        (ex_dir / "lora.txt").write_text(lora, encoding="utf-8")

    # auto model: Wan I2V studio template (or the base chosen in the bowl)
    base_key = base.strip() if (base or "").strip() in BASE_MODELS else "Wan 2.2 I2V Lightning Studio (Example-driven)"
    template_meta = BASE_MODELS[base_key]
    if template_meta.get("template") == "qwen_object_manip":
        # the manipulater's "LoRA" list is its adapter set, not the warehouse
        all_loras = [
            "Qwen-Image-Edit-2511-Object-Adder",
            "Qwen-Image-Edit-2511-Object-Remover",
            "QIE-2511-Object-Remover-v2",
            "Zoom-Master",
            "Extract-Outfit",
            "Outfit-Design-Layout",
        ]
    else:
        assets = (warehouse_list().get("assets") or [])
        all_loras = [a["name"] for a in assets] or [lora]
    cfg = {**State.config,
           "base_model": base_key, "loras": all_loras or [lora],
           "space_name": "", "steps": template_meta.get("steps", 6),
           "guidance": template_meta.get("guidance", 1.0),
           "seconds": 3.0}
    State.config = cfg

    name = f"bowl-{n_ex}-{int(time.time()) % 100000}"
    examples = examples_list().get("examples") or []
    out = _render_space(name, base_key, all_loras or [lora], cfg, examples)
    space_id = f"{OWNER}/{name}"
    api = HfApi(token=token)
    try:
        api.create_repo(repo_id=space_id, repo_type="space", private=True,
                        space_sdk="gradio", space_hardware=SpaceHardware.ZERO_A10G,
                        exist_ok=True)
    except Exception:
        try:
            api.create_repo(repo_id=space_id, repo_type="space", private=True,
                            space_sdk="gradio", exist_ok=True)
        except Exception as e2:
            raise HTTPException(status_code=400, detail=f"Cannot create Space: {e2}")
    try:
        upload_folder(repo_id=space_id, folder_path=str(out), repo_type="space",
                      token=token, commit_message="Workshop: bowl packet installed")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Upload failed: {e}")
    return {"status": "ok", "space": space_id,
            "url": f"https://huggingface.co/spaces/{space_id}",
            "examples": len(examples), "prompt": prompt, "lora": lora}


@app.post("/api/warehouse/upload")
async def warehouse_upload(file: UploadFile = File(...)):
    """Push a dropped .safetensors into the reference LoRA library."""
    if not State.vault_token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    name = file.filename or "lora.safetensors"
    if not name.endswith(".safetensors"):
        raise HTTPException(status_code=400, detail="Only .safetensors assets are accepted.")
    data = await file.read()
    tmp = Path(tempfile.gettempdir()) / f"ws_upload_{abs(hash(name))}_{name}"
    tmp.write_bytes(data)
    try:
        api = HfApi(token=State.vault_token)
        api.upload_file(
            path_or_fileobj=str(tmp), path_in_repo=name,
            repo_id=WAREHOUSE_REPO, repo_type="model",
            commit_message=f"Workshop: add LoRA {name}",
        )
    finally:
        tmp.unlink(missing_ok=True)
    return {"status": "ok", "name": name}


@app.post("/api/token/verify")
def token_verify(payload: dict):
    token = (payload.get("token") or "").strip()
    if not token:
        raise HTTPException(status_code=400, detail="Token is required.")
    try:
        api = HfApi(token=token)
        info = api.whoami()
    except Exception as e:
        raise HTTPException(status_code=401, detail="Invalid token: " + str(e))
    State.vault_token = token
    return {"status": "success", "username": info.get("name"), "auth": True}


@app.get("/api/state")
def state_get():
    return {
        "owner": OWNER,
        "vault": bool(State.vault_token),
        "username": _safe_whoami(),
        "warehouse": WAREHOUSE_REPO,
        "base_models": list(BASE_MODELS.keys()),
        "config": State.config,
    }


# ---------------------------------------------------------------------------
# Config (the pipeline being assembled)
# ---------------------------------------------------------------------------
class ConfigPayload(BaseModel):
    base_model: Optional[str] = None
    loras: Optional[list[str]] = None
    steps: Optional[int] = None
    guidance: Optional[float] = None
    fps: Optional[int] = None
    seconds: Optional[float] = None
    resolution: Optional[str] = None
    title: Optional[str] = None
    space_name: Optional[str] = None


@app.post("/api/config")
def config_save(payload: ConfigPayload):
    cfg = State.config
    for key, val in payload.model_dump(exclude_none=True).items():
        cfg[key] = val
    if not cfg.get("base_model"):
        cfg["base_model"] = list(BASE_MODELS.keys())[0]
    State.config = cfg
    return {"status": "success", "config": State.config}


# ---------------------------------------------------------------------------
# Examples (labeled image + prompt pairs, baked into published Spaces)
# ---------------------------------------------------------------------------
@app.post("/api/example")
async def example_add(files: list[UploadFile] = File(...), prompt: str = Form(""),

                      lora: str = Form("")):
    """Add an example: one or MORE images + a prompt (+ optional LoRA tag).



    Multi-image examples are how the Qwen-Image-Edit-2511 template works

    (lighting transfer, style transfer, etc.).

    """
    if not State.vault_token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    if not files:
        raise HTTPException(status_code=400, detail="Pick at least one image.")
    EXAMPLES_DIR.mkdir(parents=True, exist_ok=True)
    idx = 1
    while (EXAMPLES_DIR / f"example_{idx}").exists():
        idx += 1
    ex_dir = EXAMPLES_DIR / f"example_{idx}"
    ex_dir.mkdir()
    for i, f in enumerate(files):
        ext = Path(f.filename or "img.jpg").suffix.lower() or ".jpg"
        (ex_dir / f"image_{i + 1}{ext}").write_bytes(await f.read())
    (ex_dir / "prompt.txt").write_text((prompt or "").strip(), encoding="utf-8")
    (ex_dir / "lora.txt").write_text((lora or "").strip(), encoding="utf-8")
    return {"status": "ok", "count": idx, "images": [p.name for p in ex_dir.glob("image_*")]}


@app.get("/api/examples")
def examples_list():
    """Each example may reference multiple images."""
    if not EXAMPLES_DIR.exists():
        return {"examples": []}
    items = []
    for ex_dir in sorted(EXAMPLES_DIR.glob("example_*"), key=lambda p: p.name):
        if not ex_dir.is_dir():
            continue
        images = sorted(p.name for p in ex_dir.glob("image_*"))
        prompt_txt = ex_dir / "prompt.txt"
        lora_txt = ex_dir / "lora.txt"
        prompt = prompt_txt.read_text(encoding="utf-8") if prompt_txt.exists() else ""
        lora = lora_txt.read_text(encoding="utf-8") if lora_txt.exists() else ""
        items.append({"images": images, "prompt": prompt, "lora": lora})
    return {"examples": items}


# ---------------------------------------------------------------------------
# Template rendering + publish
# ---------------------------------------------------------------------------
class PublishPayload(BaseModel):
    space_id: Optional[str] = None
    visibility: str = "private"


def _render_space(space_name: str, base_key: str, loras: list[str],

                  cfg: dict, examples: list[dict]) -> Path:
    base = BASE_MODELS[base_key]
    tmpl_dir = Path(__file__).parent / "templates" / base.get("template", "wan22_i2v")
    out = Path(SANDBOX) / _slugify(space_name)
    if out.exists():
        shutil.rmtree(out)
    (out / "examples").mkdir(parents=True, exist_ok=True)

    title = _slugify(space_name).replace("-", " ").title()
    vars_ = {
        "PIPELINE_ID": base["pipeline_id"],
        "BASE_MODEL": base["repo_id"],
        "STEPS_DEFAULT": str(cfg.get("steps", base.get("steps", 6))),
        "GUIDANCE_DEFAULT": str(cfg.get("guidance", base.get("guidance", 1.0))),
        "DURATION_DEFAULT": str(cfg.get("seconds", 3.0)),
        "MAX_FRAMES": str(base.get("max_frames", 241)),
        "LORA_REPO": WAREHOUSE_REPO,
        "LORA_CHOICES": json.dumps(loras),
        "SPACE_NAME": space_name,
        "TITLE": title,
    }

    # Copy static template tree first (qwenimage/, index.html, etc.)
    for item in tmpl_dir.iterdir():
        if item.is_dir():
            shutil.copytree(item, out / item.name, dirs_exist_ok=True)
        elif item.name not in ("app.py", "requirements.txt", "README.md"):
            shutil.copy(item, out / item.name)

    if base.get("multi_image"):
        # Qwen-Image-Edit style: EXAMPLES_CONFIG entries carry image LISTS + a LoRA tag.
        entries = []
        for ex in examples:
            imgs = ex.get("images") or []
            entry = {
                "images": [f"examples/{n}" for n in imgs],
                "prompt": (ex.get("prompt") or "")[:160],
            }
            if ex.get("lora"):
                entry["lora"] = ex["lora"]
            entries.append(entry)
        vars_["EXAMPLES_CONFIG"] = "[\n    " + ",\n    ".join(json.dumps(e) for e in entries) + "\n]"
        vars_["EXAMPLES"] = ""
        rels = ("app.py", "requirements.txt", "pre-requirements.txt", "README.md", "index.html")
    else:
        vars_["EXAMPLES_CONFIG"] = ""
        vars_["EXAMPLES"] = ",\n        ".join(
            f"[{json.dumps((ex.get('prompt') or '')[:120])}, 'examples/{ex['images'][0]}']"
            for ex in examples if ex.get("images")
        ) or f"[{json.dumps('a gentle cinematic shot, warm light')}, 'examples/example_1.jpg']"
        rels = ("app.py", "requirements.txt", "README.md")

    for rel in rels:
        src = tmpl_dir / rel
        if not src.exists():
            continue
        text = src.read_text(encoding="utf-8")
        for k, v in vars_.items():
            text = text.replace("{{" + k + "}}", str(v))
        (out / rel).write_text(text, encoding="utf-8")

    for ex in examples:
        for img in ex.get("images") or []:
            src_img = Path(EXAMPLES_DIR) / "example_*" / img
            for match in sorted(Path(EXAMPLES_DIR).glob("example_*")):
                cand = match / img
                if cand.exists():
                    shutil.copy(cand, out / "examples" / img)
    # Guarantee at least one example image exists so gr.Examples never dangles.
    if not any((out / "examples").iterdir()):
        placeholder = out / "examples" / "example_1.jpg"
        placeholder.write_bytes(
            bytes.fromhex(
                "FFD8FFE000104A46494600010100000100010000FFDB004300080606070605080707070909080A0C140D0C0B0B0C1912130F141D1A1F1E1D1A1C1C20242E2720222C231C1C2837292C30313434341F27393D38323C2E333432FFC0000B080040004001011100FFC4001F0000010501010101010100000000000000000102030405060708090A0BFFC400B5100002010303020403050504040000017D01020300041105122131410613516107227114328191A1082342B1C11552D1F02433627282090A161718191A25262728292A3435363738393A434445464748494A535455565758595A636465666768696A737475767778797A838485868788898A92939495969798999AA2A3A4A5A6A7A8A9AAB2B3B4B5B6B7B8B9BAC2C3C4C5C6C7C8C9CAD2D3D4D5D6D7D8D9DAE1E2E3E4E5E6E7E8E9EAF1F2F3F4F5F6F7F8F9FAFFC4001F0100030101010101010101010000000000000102030405060708090A0BFFC400B51100020102040403040705040400010277000102031104052131061241510761711322328108144291A1B1C109233352F0156272D10A162434E125F11718191A262728292A35363738393A434445464748494A535455565758595A636465666768696A737475767778797A82838485868788898A92939495969798999AA2A3A4A5A6A7A8A9AAB2B3B4B5B6B7B8B9BAC2C3C4C5C6C7C8C9CAD2D3D4D5D6D7D8D9DAE2E3E4E5E6E7E8E9EAF2F3F4F5F6F7F8F9FAFFDA000C03010002110311003F000000FFFF",
            )
        )
    return out


@app.post("/api/publish")
def publish(payload: PublishPayload):
    token = _token()
    if not token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    cfg = State.config
    if not (payload.space_id or cfg.get("space_name")):
        raise HTTPException(status_code=400, detail="No Space name set.")

    space_name = (payload.space_id or str(cfg["space_name"])).strip()
    space_id = f"{OWNER}/{space_name}" if "/" not in space_name else space_name
    base_key = cfg.get("base_model") or list(BASE_MODELS.keys())[0]
    if base_key not in BASE_MODELS:
        base_key = list(BASE_MODELS.keys())[0]
    loras = cfg.get("loras") or []
    examples = examples_list().get("examples", []) or []

    out = _render_space(space_name, base_key, loras, cfg, examples)

    api = HfApi(token=token)
    try:
        api.create_repo(
            repo_id=space_id, repo_type="space", private=True,
            space_sdk="gradio", space_hardware=SpaceHardware.ZERO_A10G,
            exist_ok=True,
        )
    except Exception:
        try:
            api.create_repo(
                repo_id=space_id, repo_type="space", private=True,
                space_sdk="gradio", exist_ok=True,
            )
        except Exception as e2:
            raise HTTPException(status_code=400, detail=f"Cannot create Space: {e2}")

    try:
        upload_folder(
            repo_id=space_id, folder_path=str(out), repo_type="space",
            token=token, commit_message="Workshop: generated LoRA pipeline",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Upload failed: {e}")

    if payload.visibility == "public":
        try:
            api.set_space_visibility(space_id, public=True)
        except Exception:
            pass

    return {"status": "ok", "space": space_id, "url": f"https://huggingface.co/spaces/{space_id}"}


@app.get("/api/spaces")
def spaces_list():
    if not State.vault_token:
        return {"spaces": [], "error": "vault empty"}
    try:
        api = HfApi(token=State.vault_token)
        out = []
        for s in api.list_spaces(author=OWNER):
            try:
                runtime = api.get_space_runtime(s.id).stage
            except Exception:
                runtime = "unknown"
            out.append({"id": s.id, "runtime": runtime})
        return {"spaces": out}
    except Exception as e:
        return {"spaces": [], "error": str(e)}


@app.get("/api/space_logs")
def space_logs(space: str = ""):
    """Debug: raw runtime info + error detail for one of the owner's Spaces."""
    if not State.vault_token:
        return {"error": "vault empty"}
    if not space:
        return {"error": "space id required"}
    try:
        api = HfApi(token=State.vault_token)
        rt = api.get_space_runtime(space)
        info = {"stage": getattr(rt, "stage", None), "hardware": str(getattr(rt, "hardware", "")),
                "runtime": getattr(rt, "runtime", None)}
        for cand in ("error", "last_message", "message", "error_message"):
            try:
                v = getattr(rt, cand)
                if v:
                    info[cand] = str(v)[:800]
            except Exception:
                pass
        logs = api.get_space_logs(space) if hasattr(api, "get_space_logs") else None
        tail = ""
        if logs is not None:
            try:
                if hasattr(logs, "logs"):
                    tail = (logs.logs or "")[-4000:]
                else:
                    tail = str(logs)[-4000:]
            except Exception as e:
                tail = f"[logs err: {e}]"
        return {"info": info, "log_tail": tail}
    except Exception as e:
        return {"error": str(e)}


class VisPayload(BaseModel):
    space: str
    public: bool = False


# ---------------------------------------------------------------------------
# URL Quick flow: paste a URL -> find the same working model + examples,
# install the whole packet as one Space (2-window flow)
# ---------------------------------------------------------------------------
_IMG_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp")


class UrlFlowPayload(BaseModel):
    url: str = ""
    mode: str = "find"  # find | install
    space_name: Optional[str] = None
    visibility: str = "private"


def _resolve_hf_id(url: str) -> Optional[str]:
    url = (url or "").strip()
    if not url:
        return None
    if url.count("/") == 1 and not url.startswith(("http://", "https://")):
        return url
    m = re.search(r"huggingface\.co/(?:spaces|models)/([A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)", url)
    return m.group(1) if m else None


def _fetch_remote_text(api: HfApi, repo_id: str, repo_type: str) -> tuple:
    """Fetch (app.py text, repo files) from a remote Space or model repo."""
    files = set()
    try:
        files = set(api.list_repo_files(repo_id, repo_type=repo_type))
    except Exception:
        pass
    text = ""
    for cand in ("app.py", "main.py", "README.md"):
        if cand in files:
            try:
                text += hf_hub_download(repo_id, cand, repo_type=repo_type,
                                        token=State.vault_token) + "\n"
                with open(hf_hub_download(repo_id, cand, repo_type=repo_type,
                                          token=State.vault_token), encoding="utf-8",
                          errors="replace") as fh:
                    text += fh.read()
            except Exception:
                continue
    return text, sorted(files)


def _extract_examples_cfg(text: str) -> list[dict]:
    m = re.search(r"EXAMPLES_CONFIG\s*=\s*(\[[\s\S]*?\n\])", text)
    if not m:
        m = re.search(r"EXAMPLES_CONFIG\s*=\s*(\[[\s\S]*?\])", text)
    if not m:
        return []
    try:
        cfg = json.loads(m.group(1))
    except Exception:
        return []
    return [e for e in cfg if isinstance(e, dict) and
            (e.get("images") or e.get("prompt") or e.get("lora"))]


@app.post("/api/urlflow")
def url_flow(payload: UrlFlowPayload):
    token = _token()
    if not token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    repo_id = _resolve_hf_id(payload.url or "")
    if not repo_id:
        raise HTTPException(status_code=400,
                            detail="Paste a URL like https://huggingface.co/spaces/org/name (or org/name).")

    api = HfApi(token=token)
    info = None
    try:
        info = api.space_info(repo_id)
        repo_type = "space"
    except Exception:
        try:
            api.model_info(repo_id)
            repo_type = "model"
        except Exception:
            raise HTTPException(status_code=404,
                                detail=f"Cannot resolve {repo_id} - not a Space or model repo.")

    text, files = _fetch_remote_text(api, repo_id, repo_type)
    sig = _analyze_app(text) if text else {}
    found_model = None
    for key, base in BASE_MODELS.items():
        if (sig.get("base_model")) and (sig["base_model"] in base["repo_id"] or base["repo_id"] in sig["base_model"]):
            found_model = key
            break
    if not found_model and sig.get("pipeline"):
        for key, base in BASE_MODELS.items():
            if base["pipeline_id"] == sig["pipeline"]:
                found_model = key
                break

    # find examples (app config first, then images in repo)
    ex_cfg = _extract_examples_cfg(text)
    remote_imgs = [f for f in files if any(f.lower().endswith(e) for e in _IMG_EXTS)]
    if not ex_cfg and remote_imgs:
        for img in remote_imgs[:4]:
            ex_cfg.append({"images": [img], "prompt": "", "lora": ""})

    answer = {
        "repo": repo_id, "type": repo_type,
        "found_model": found_model,
        "workflow": {"pipeline": sig.get("pipeline"), "steps": sig.get("steps"),
                     "guidance": sig.get("guidance"), "fps": sig.get("fps"),
                     "loras": sig.get("loras"), "quantization": sig.get("quantization")},
        "examples": ex_cfg,
        "status": "ok",
    }
    if payload.mode == "install":
        return _install_packet(payload, repo_id, repo_type, ex_cfg, found_model, answer)
    return answer


def _install_packet(payload, repo_id, repo_type, ex_cfg, found_model, answer) -> dict:
    api = HfApi(token=_token())
    if not found_model:
        raise HTTPException(status_code=400,
                            detail="No matching recipe found for that model - stop, check the URL, or pick a model manually.")
    # pull every example image into the local sandbox as a whole packet
    imports = []
    for i, ex in enumerate(ex_cfg):
        imgs_rel = []
        for img in (ex.get("images") or []):
            if img.startswith(("http://", "https://", "data:")):
                imgs_rel.append(img)
                continue
            try:
                local = hf_hub_download(repo_id, img, repo_type=repo_type,
                                        token=_token())
            except Exception:
                continue
            dst = Path(EXAMPLES_DIR) / f"example_{100 + i}" / f"image_{len(imgs_rel)+1}{os.path.splitext(img)[1] or '.jpg'}"
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(local, dst)
            imgs_rel.append(dst.name)
        if not imgs_rel:
            continue
        ex_dir = Path(EXAMPLES_DIR) / f"example_{100 + i}"
        ex_dir.mkdir(parents=True, exist_ok=True)
        (ex_dir / "prompt.txt").write_text((ex.get("prompt") or "")[:220], encoding="utf-8")
        if ex.get("lora"):
            (ex_dir / "lora.txt").write_text(ex["lora"], encoding="utf-8")
        imports.append({"images": sorted(p.name for p in ex_dir.glob("image_*")),
                        "prompt": ex.get("prompt", "")[:220], "lora": ex.get("lora", "")})
    imported = imports or []
    if not imported:
        raise HTTPException(status_code=400,
                            detail="No importable example images found at that URL.")

    # grab the LoRA choices from the warehouse
    assets = warehouse_list().get("assets", []) or []
    loras = [a["name"] for a in assets]

    if not payload.space_name:
        slug = _slugify(repo_id.split("/")[-1]).replace("-", " ")
        cfg = {**State.config, "base_model": found_model, "loras": loras}
    else:
        cfg = {**State.config, "base_model": found_model, "space_name": payload.space_name.strip(), "loras": loras}
    space_name = (cfg.get("space_name") or slug).strip()
    cfg["space_name"] = space_name

    out = _render_space(space_name, found_model, loras, cfg, imported)
    space_id = f"{OWNER}/{space_name}" if "/" not in space_name else space_name
    try:
        api.create_repo(repo_id=space_id, repo_type="space", private=True,
                        space_sdk="gradio", space_hardware=SpaceHardware.ZERO_A10G,
                        exist_ok=True)
    except Exception:
        try:
            api.create_repo(repo_id=space_id, repo_type="space", private=True,
                            space_sdk="gradio", exist_ok=True)
        except Exception as e2:
            raise HTTPException(status_code=400, detail=f"Cannot create Space: {e2}")
    try:
        upload_folder(repo_id=space_id, folder_path=str(out), repo_type="space",
                      token=_token(), commit_message="Workshop: whole packet install")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Upload failed: {e}")
    if payload.visibility == "public":
        try:
            api.set_space_visibility(space_id, public=True)
        except Exception:
            pass
    answer["installed"] = True
    answer["packet_space"] = space_id
    answer["url"] = f"https://huggingface.co/spaces/{space_id}"
    answer["imported_examples"] = imported
    return answer


@app.post("/api/visibility")
def set_visibility(payload: VisPayload):
    token = _token()
    if not token:
        raise HTTPException(status_code=403, detail="Security vault empty.")
    try:
        HfApi(token=token).set_space_visibility(payload.space, public=bool(payload.public))
        return {"status": "ok", "space": payload.space, "public": bool(payload.public)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# Space research (find + compare similar working models)
# ---------------------------------------------------------------------------
SEARCH_PARAMS = {  # queried via the HF Hub API search endpoint
    "limit": 12,
}


def hf_search_spaces(query: str) -> list[dict]:
    """Search for Spaces on the Hub that match a free-form query."""
    if not State.vault_token:
        return []
    try:
        api = HfApi(token=State.vault_token)
        out = []
        for s in api.list_spaces(search=query, limit=SEARCH_PARAMS["limit"]):
            try:
                info = api.space_info(s.id)
            except Exception:
                info = None
            out.append({
                "id": s.id,
                "likes": int(getattr(info, "likes", 0) or 0) if info else 0,
                "tags": list(getattr(info, "tags", []) or []) if info else [],
                "sdk": getattr(info, "sdk", "") if info else "",
                "runtime": getattr(getattr(info, "runtime", None), "stage", "") if info else "",
            })
        return out
    except Exception as e:
        return [{"error": str(e)}]


def _space_main_file(api: HfApi, space_id: str) -> Optional[str]:
    """Find the primary .py file of a Space (app.py first, then main.py)."""
    for cand in ("app.py", "main.py"):
        try:
            if cand in api.list_repo_files(space_id, repo_type="space"):
                return hf_hub_download(space_id, cand, repo_type="space",
                                       token=State.vault_token)
        except Exception:
            continue
    return None


def _analyze_app(text: str) -> dict:
    """Extract the workflow signature of a Gradio app.py."""
    import ast as _ast
    sig = {
        "pipeline": None, "base_model": None, "steps": None, "guidance": None,
        "fps": None, "max_frames": None, "loras": bool(re.search(r"load_lora_weights|LoRA|lora", text, re.I)),
        "quantization": None, "examples": False, "i2v": bool(re.search(r"Image.*Video|wan.*i2v|image", text.lower())),
    }
    # static pipeline assignment
    m = re.search(r"(WanImageToVideoPipeline|WanVideoToVideoPipeline|WanTextToVideoPipeline|StableDiffusionPipeline|sd3|FLUX|SDXLPipeline|DiffusionPipeline)", text)
    if m:
        sig["pipeline"] = m.group(1)
    # base model repo ids are the most reliable cross-reference
    ids = []
    for pat in (r"from_pretrained\(\s*[\"']([^\"']+)[\"']",
                r"(?:BASE_MODEL|base_model_id|pretrained_model_name_or_path)\s*[=:]\s*[\"']([^\"']+)[\"']",
                r"[\"']([^\"']+)[\"']\s*,?\s*$"):
        ids = re.findall(pat, text, re.I | re.M)
        if ids:
            break
    if not ids:
        ids = re.findall(r"[\"']([^\"']*[Ww]an[^\"']*)[\"']", text)
    ids = [i for i in ids if "/" in i]
    sig["base_model"] = ids[0] if ids else None
    m = re.search(r"(?:num_inference_steps|STEPS_DEFAULT)\s*[=:]\s*(\d+)", text)
    if m:
        sig["steps"] = int(m.group(1))
    m = re.search(r"guidance_scale\s*[=:]\s*([\d.]+)", text)
    if m:
        sig["guidance"] = float(m.group(1))
    m = re.search(r"fps\s*[=:]\s*(\d+)", text)
    if m:
        sig["fps"] = int(m.group(1))
    m = re.search(r"(?:MAX_FRAMES|max_frames)\s*[=:]\s*(\d+)", text)
    if m:
        sig["max_frames"] = int(m.group(1))
    if re.search(r"quantize_|torchao|FP8|Float8", text):
        sig["quantization"] = True
    if re.search(r"gr\.Examples|examples\s*=", text):
        sig["examples"] = True
    return sig


def _compare_workflow(remote: dict, local_cfg: dict) -> dict:
    """Cross-reference the remote Space app.py workflow with our configuration."""
    deltas = []
    base_key = ours_base = BASE_MODELS.get(local_cfg.get("base_model") or list(BASE_MODELS)[0], {})
    if remote.get("pipeline") and "WanImageToVideoPipeline" not in remote["pipeline"]:
        deltas.append(f"their pipeline is {remote['pipeline']} (ours is WanImageToVideoPipeline)")
    if remote.get("base_model") and ours_base.get("repo_id") not in (remote["base_model"] or ""):
        deltas.append(f"their base model id differs: {remote['base_model'][:60]}")
    if remote.get("steps") is not None and remote["steps"] != ours_base.get("steps"):
        deltas.append(f"steps: theirs {remote['steps']} vs ours {ours_base.get('steps')}")
    if remote.get("guidance") is not None and abs(remote["guidance"] - ours_base.get("guidance", 1.0)) > 0.1:
        deltas.append(f"guidance: theirs {remote['guidance']} vs ours {ours_base.get('guidance')}")
    if remote.get("fps") and remote["fps"] != 16:
        deltas.append(f"fps: theirs {remote['fps']} vs ours 16")
    return {
        "theirs": {"steps": remote.get("steps"), "guidance": remote.get("guidance"),
                   "fps": remote.get("fps"), "pipeline": remote.get("pipeline"),
                   "base_model": remote.get("base_model"), "loras": remote.get("loras"),
                   "quantization": remote.get("quantization"), "examples": remote.get("examples")},
        "deltas": deltas,
    }


def _fetch_space(space_id: str) -> tuple:
    """Return (space_info, app.py path, text) for a Space."""
    api = HfApi(token=State.vault_token)
    try:
        info = api.space_info(space_id).__dict__
    except Exception:
        info = {}
    path = _space_main_file(api, space_id)
    text = ""
    if path:
        with open(path, encoding="utf-8", errors="replace") as fh:
            text = fh.read()
    return info, path, text


# ---------------------------------------------------------------------------
# Diagnosis rubric: what "good" looks like for a Wan I2V LoRA Studio Space
# ---------------------------------------------------------------------------
# The gold standard is the owner's proven reference Space
# (wan2-2-fp8da-aoti-preview-nsfw) plus diffusers/runtime best practices.
# Every check is statically verifiable from app.py + repo metadata, so the
# assistant can call good/bad without guessing.
RUBRIC = [
    {"id": "base_exists", "name": "Base model repo resolves on the Hub",
     "missing": "base model id could not be verified to exist - check the repo id",
     "warn": "", "cat": "core"},
    {"id": "pipeline_class", "name": "Uses the correct Wan I2V diffusers pipeline",
     "missing": "no WanImageToVideoPipeline import/use detected", "warn": "", "cat": "core"},
    {"id": "to_cuda_dtype", "name": "Sets torch_dtype (bf16/fp16) on load",
     "missing": "no torch_dtype on from_pretrained - will default to fp32 and likely OOM",
     "warn": "", "cat": "core"},
    {"id": "quantization", "name": "Quantizes for VRAM (FP8 on Hopper+, INT8 fallback)",
     "missing": "no quantize_ / torchao - model runs bf16 and likely OOMs on <24GB cards",
     "warn": "quantize_ used but no INT8 fallback for <Hopper GPUs", "cat": "core"},
    {"id": "vram_cleanup", "name": "Frees VRAM between generations (gc + empty_cache)",
     "missing": "no gc.collect()/empty_cache - runs will leak VRAM and get flaky", "warn": "", "cat": "core"},
    {"id": "resolution", "name": "Resizes input to a GPU-friendly, /16 frame",
     "missing": "image passed straight through (any size) - hurts quality & VRAM", "warn": "", "cat": "core"},
    {"id": "seed_control", "name": "Supports manual + random seed",
     "missing": "no manual_seed / random seed path", "warn": "", "cat": "core"},
    {"id": "num_frames", "name": "Frame math is valid (>=9, matches a 4r+1 pattern, respects max)",
     "missing": "num_frames not clamped - GPU crash risk", "warn": "", "cat": "core"},
    {"id": "lora_handling", "name": "LoRA downloads securely (HF_TOKEN) & unloads when deselected",
     "missing": "no LoRA support OR loads without the env token",
     "warn": "downloads without token (private warehouse will 401); or no unload path", "cat": "core"},
    {"id": "examples_valid", "name": "gr.Examples point at images that actually exist",
     "missing": "no gr.Examples", "warn": "Examples reference image files that are missing in the repo",
     "cat": "ux"},
    {"id": "negative_prompt", "name": "Provides a sane default negative prompt",
     "missing": "no negative prompt input/default", "warn": "", "cat": "ux"},
    {"id": "advanced_panel", "name": "Exposes steps/guidance/seed controls",
     "missing": "no controls for steps/guidance", "warn": "", "cat": "ux"},
    {"id": "gpu_decorator", "name": "Uses @spaces.GPU for queueing/GPU allocation",
     "missing": "no spaces.GPU decorator - Space may run on CPU or time out", "warn": "", "cat": "deploy"},
    {"id": "secret_safety", "name": "Never hardcodes HF_TOKEN; reads it from env",
     "missing": "hardcoded token literal found - REVOKE it", "warn": "", "cat": "security"},
]


def _rubric_grade(checks: list[dict]) -> str:
    ok = sum(1 for c in checks if c["status"] == "ok")
    warnings_ = sum(1 for c in checks if c["status"] == "warn")
    missing = sum(1 for c in checks if c["status"] == "missing")
    total = len(checks)
    if missing == 0 and warnings_ == 0 and ok == total:
        return "Perfect"
    if missing == 0:
        return "Solid"
    bad_core = sum(1 for c in checks if c["status"] == "missing" and c.get("cat") not in ("ux", "deploy"))
    if bad_core:
        return "Unreliable"
    return "Needs polish"


def _diagnose_space(space_id: str) -> dict:
    """Grade a Space against the rubric; propose add-ons afterwards."""
    api = HfApi(token=State.vault_token)
    info, path, text = _fetch_space(space_id)
    out: dict = {"space": space_id,
                 "likes": int(info.get("likes") or 0),
                 "found": bool(text), "checks": [], "addons": []}
    if not text:
        out["error"] = f"could not read app.py from {space_id} (private or not a python Space)"
        out["grade"] = "Unknown"
        return out

    files = []
    try:
        files = api.list_repo_files(space_id, repo_type="space")
    except Exception:
        pass

    base_ids = re.findall(r"from_pretrained\(\s*[\"']([^\"']+)[\"']", text)
    base_id = base_ids[0] if base_ids else None
    base_ok = bool(base_id)
    if base_id:
        try:
            api.model_info(base_id)
        except Exception:
            base_ok = False

    def _check(cid, status, evidence, cat="core"):  # noqa
        out["checks"].append({"id": cid, "status": status, "evidence": evidence, "cat": cat})

    _check("base_exists", "ok" if base_ok else "missing",
           f"resolves on hub: {bool(base_ok)}" + (f" ({base_id})" if base_id else ""))
    pipe = "WanImageToVideoPipeline" if re.search(r"WanImageToVideoPipeline", text) else (
        "WanVideoToVideoPipeline" if re.search(r"WanVideoToVideoPipeline", text) else None)
    _check("pipeline_class", "ok" if pipe else "missing",
           pipe or "no Wan diffusers I2V pipeline")
    dt = bool(re.search(r"torch_dtype\s*=\s*torch\.(bfloat16|float16)", text))
    _check("to_cuda_dtype", "ok" if dt else "missing",
           "bf16/fp16 on load" if dt else "no dtype set - likely fp32/OOM")
    quant_has = bool(re.search(r"quantize_|Float8|Int8WeightOnly", text))
    quant_fallback = bool(re.search(r"Int8WeightOnlyConfig|has_fp8|device_capability", text)) if quant_has else False
    _check("quantization", "ok" if quant_has and quant_fallback else ("warn" if quant_has else "missing"),
           "FP8+INT8 fallback" if quant_has and quant_fallback else ("quantized, no fallback" if quant_has else "no quantization"))
    _check("vram_cleanup", "ok" if re.search(r"empty_cache", text) else "missing",
           "clears VRAM between runs" if re.search(r"empty_cache", text) else "no gc/empty_cache")
    _check("resolution", "ok" if re.search(r"resize|multiple_of|LANCZOS", text) else "missing",
           "resizes to /16 frame" if re.search(r"multiple_of|resize", text) else "no resize handling")
    _check("seed_control", "ok" if re.search(r"manual_seed|generator", text) else "missing",
           "seeded generator" if re.search(r"manual_seed", text) else "no seed control")
    frames_ok = bool(re.search(r"num_frames.*(?:np\.clip|clip\(|%|//)|MAX_FRAMES", text, re.S))
    _check("num_frames", "ok" if frames_ok else "warn",
           "clamped frame count" if frames_ok else "frames not clamped - crash risk")
    lora_has = bool(re.search(r"load_lora_weights", text))
    lora_token = bool(re.search(r"hf_hub_download.*HF_TOKEN|HF_TOKEN", text))
    _check("lora_handling", "ok" if lora_has and lora_token else ("warn" if lora_has else "missing"),
           ("LoRA + env token" if lora_has and lora_token else
            ("LoRA but no token flow" if lora_has else "no LoRA support")))
    ex_images = re.findall(r"[\"']examples/([^\"']+)[\"']", text)
    ex_missing = [e for e in ex_images if e not in files]
    _check("examples_valid", "ok" if ex_images and not ex_missing else ("warn" if ex_images else "missing"),
           (f"{len(ex_images)} example image(s)" if ex_images and not ex_missing else
            (f"missing images: {ex_missing[:3]}" if ex_missing else "no gr.Examples")),
           cat="ux")
    _check("negative_prompt", "ok" if re.search(r"negative_prompt", text) else "missing",
           "negative prompt present" if re.search(r"negative_prompt", text) else "no negative prompt",
           cat="ux")
    _check("advanced_panel", "ok" if re.search(r"guidance_scale|num_inference_steps|Slider", text) else "missing",
           "steps/guidance controls" if re.search(r"guidance_scale|num_inference_steps|Slider", text) else "no advanced controls",
           cat="ux")
    _check("gpu_decorator", "ok" if re.search(r"@spaces\.GPU|spaces\.GPU", text) else "missing",
           "GPU queue decorator" if re.search(r"spaces\.GPU", text) else "no @spaces.GPU",
           cat="deploy")
    secret_leak = re.search(r"hf_[A-Za-z0-9]{10,}", text)
    _check("secret_safety", "ok" if not secret_leak else "missing",
           "token read from env" if not secret_leak else "HARDCODED TOKEN FOUND - revoke it")

    out["grade"] = _rubric_grade(out["checks"])

    out["addons"] = []
    if not re.search(r"SageAttention", text):
        out["addons"].append("SageAttention for faster attention on modern GPUs")
    if not re.search(r"torch\.compile|dynamo", text):
        out["addons"].append("torch.compile / dynamo over the transformer for speed")
    if not re.search(r"cache_examples=True", text):
        out["addons"].append("cache_examples=True so preset runs render instantly")
    if not re.search(r"seed.*out|Seed", text, re.I):
        out["addons"].append("return + display the used seed so runs are reproducible")
    if not re.search(r"prompt.*enhance|enhance.*prompt", text, re.I):
        out["addons"].append("optional prompt enhancer (LLM upscale) before inference")
    return out


def _diagnose_mine() -> dict:
    """Run the rubric over the owner's own published Spaces, best-first."""
    api = HfApi(token=State.vault_token)
    results = []
    try:
        mine = [s.id for s in api.list_spaces(author=OWNER)]
    except Exception:
        return {"spaces": [], "error": "cannot list your Spaces"}
    for sid in mine:
        try:
            results.append(_diagnose_space(sid))
        except Exception as e:
            results.append({"space": sid, "error": str(e), "grade": "Unknown"})
    order = {"Unreliable": 0, "Needs polish": 1, "Solid": 2, "Perfect": 3, "Unknown": 4}
    results.sort(key=lambda d: (order.get(d.get("grade"), 4), -(d.get("likes") or 0)))
    return {"spaces": results, "count": len(results)}


# ---------------------------------------------------------------------------
# Chat assistant
# ---------------------------------------------------------------------------
class ChatMsg(BaseModel):
    messages: list[dict]


SYSTEM_PROMPT = """You are the Workshop Build Assistant for a LoRA pipeline studio owned by {owner}.

You talk like a collaborator with a friendly, technical tone - not like a support bot.



Context you can rely on:

- Reference LoRA library (warehouse): {warehouse}

- You configure a Wan 2.2 I2V base pipeline, attach LoRA .safetensors by name, label example image+prompt pairs, then publish a private Space.

- Everything you do happens through tool calls embedded in your reply.



When you want the app to do something, output a single JSON object at the END of your reply with this exact shape:

{{"tool": "<name>", "args": {{...}}}}



Supported tools:

- "list": no args. Refresh pipeline config + warehouse assets, summarize briefly.

- "set_config": {{"base_model": "Wan 2.2 I2V Lightning 14B (AoT FP8)", "steps": 6, "guidance": 1.0, "seconds": 3.0, "space_name": "my-space", "loras": ["WAN-2.2-I2V-FaceDownAssUp-HIGH-v1.safetensors"]}}

  Configure the pipeline. LoRA names must match warehouse assets exactly.

- "publish": {{"visibility": "private"}}  Publish the configured pipeline as a private Space under {owner}. Keep private unless the owner explicitly asks for public.

- "visibility": {{"space": "<id>", "public": false}}  Flip a Space's visibility.

- "examples": {{"prompt": "the prompt for a labeled example"}}  Record a prompt; images are uploaded separately in the UI.

- "find": {{"query": "wan image to video lora space"}}  Search the Hub for similar working Spaces. Returns id, likes, tags, sdk, runtime. Use this whenever the owner wants to find similar models or compare with other workflows.

- "compare": {{"space": "<space id>"}}  Fetch that Space's app.py, extract its workflow (pipeline, base model, steps, guidance, fps, frames, LoRA use, quantization, examples) and cross-reference it with the owner's configured workflow. Reports deltas, e.g. different steps or a base model that is more popular. Ask the owner whether to refigure the config to match or paste the app.py for inspection (the tool result already includes a summary and key source lines).

- "diagnose": {{"space": "<space id>"}}  Grade a Space against a 14-point gold-standard rubric (base resolves, right pipeline, dtype, quantization with fallback, VRAM cleanup, resolution, seed control, frame math, LoRA+token flow, valid examples, negative prompt, controls, @spaces.GPU, no leaked secrets). Returns grade (Perfect/Solid/Needs polish/Unreliable), the per-check pass/fail list, and add-on suggestions. Use it AFTER find/compare to say what is genuinely wrong before proposing matches.

- "diagnose_mine": no args. Usually the first thing to check when the owner says things are "broken" or "unreliable" - grade every Space the owner has published against the same rubric, worst first, so we fix the most broken one first.



Rules:

- Respond conversationally first; only emit a tool call when the user's request maps cleanly to one.

- Keep config summaries to one or two lines.

- Be concise but human. No markdown tables.

- When comparing workflows, be specific: name the deltas (steps, guidance, base model id, likes count). Never claim theirs is better without evidence; offer to adopt the differing settings or to paste their app.py source so the owner can look at it themselves."""


def _chat_llm(messages: list[dict]) -> Optional[str]:
    headers = {"Authorization": f"Bearer {_token()}"}
    body = {"messages": messages, "max_tokens": 700, "temperature": 0.7}
    for base in CHAT_MODELS:
        url = f"https://api-inference.huggingface.co/models/{base}"
        try:
            r = requests.post(url, json=body, headers=headers, timeout=CHAT_TIMEOUT)
            if r.status_code == 200:
                data = r.json()
                choice = data.get("choices") or []
                if choice:
                    return choice[0]["message"]["content"]
        except Exception:
            continue
    return None


def _find_json_tool(text: str) -> Optional[dict]:
    idx = text.rfind("{")
    if idx < 0:
        return None
    frag = text[idx:]
    for end in range(len(frag), 0, -1):
        try:
            parsed = json.loads(frag[:end])
            if isinstance(parsed, dict) and "tool" in parsed:
                return parsed
        except Exception:
            continue
    return None


def _execute_tool(call: dict) -> str:
    name = call.get("tool")
    args = call.get("args") or {}
    try:
        if name == "list":
            w = warehouse_list()
            cfg = State.config
            return (f"Warehouse ({WAREHOUSE_REPO}): {len(w.get('assets') or [])} LoRAs. "
                    f"Config: base={cfg.get('base_model') or 'unset'}, "
                    f"loras={', '.join(cfg.get('loras') or []) or 'none'}, "
                    f"space={cfg.get('space_name') or 'unset'}.")
        if name in ("set_config", "set_space"):
            config_save(ConfigPayload(**args))
            cfg = State.config
            return "Config updated: " + " | ".join(f"{k}={v}" for k, v in cfg.items())
        if name == "publish":
            resp = publish(PublishPayload(visibility=args.get("visibility", "private")))
            return f"Published: {resp['url']} (private)"
        if name == "visibility":
            set_visibility(VisPayload(**args))
            return "Visibility updated."
        if name == "examples":
            return "Prompt noted - upload an example image in the panel to complete the pair."
        if name == "find":
            q = args.get("query") or "wan image to video"
            res = hf_search_spaces(q)
            if not res:
                return "No matching Spaces found."
            lines = ["Similar Spaces found:"]
            for r in res[:10]:
                if "error" in r:
                    continue
                lines.append(f"- {r['id']} ({r['likes']} likes, {r['sdk']}, {r['runtime']})")
            return "\n".join(lines)
        if name == "compare":
            space = args.get("space") or ""
            if not space:
                return "Specify a Space id to compare."
            info, path, text = _fetch_space(space)
            if not text:
                return f"Could not read app.py from {space} (private or no python file)."
            sig = _analyze_app(text)
            cmp = _compare_workflow(sig, State.config)
            lines = [f"Workflow of {space} (likes: {info.get('likes', 0)}):",
                     f"- pipeline: {sig['pipeline'] or 'unknown'}, base: {sig['base_model'] or 'unknown'}",
                     f"- steps={sig.get('steps')} guidance={sig.get('guidance')} fps={sig.get('fps')} max_frames={sig.get('max_frames')}"
                     + (" | quantization" if sig.get("quantization") else "")
                     + (" | LoRA" if sig.get("loras") else "")]
            if cmp["deltas"]:
                lines.append("Differences vs your config:")
                lines += [f"- {d}" for d in cmp["deltas"]]
                lines.append("Want me to refigure your config to match, or want their app.py pasted here to look at?")
            else:
                lines.append("This workflow matches your config - nothing to change.")
            return "\n".join(lines)
        if name == "diagnose":
            space = args.get("space") or ""
            if not space:
                return "Specify a Space id to diagnose."
            d = _diagnose_space(space)
            if d.get("error"):
                return d["error"]
            lines = [f"Diagnosis of {space} ({d['likes']} likes): GRADE = {d['grade']}"]
            for c in d["checks"]:
                mark = {"ok": "[ok]", "warn": "[~]", "missing": "[!!]"}.get(c["status"], "?")
                lines.append(f"{mark} {c['id']}: {c['evidence']}")
            bad = [c for c in d["checks"] if c["status"] != "ok"]
            if bad:
                lines.append(f"Fixing priority: " + ", ".join(c["id"] for c in bad[:5]) + ".")
            else:
                lines.append("Nothing broken - this follows the gold standard.")
            if d["addons"]:
                lines.append("Possible add-ons: " + "; ".join(d["addons"]) + ".")
            return "\n".join(lines)
        if name == "diagnose_mine":
            res = _diagnose_mine()
            if res.get("error"):
                return res["error"]
            spaces = res.get("spaces") or []
            if not spaces:
                return "No published Spaces found for the owner."
            lines = [f"Diagnosed {res['count']} of your Spaces, worst first:"]
            for d in spaces:
                if d.get("error"):
                    lines.append(f"- {d['space']}: {d['error']}")
                else:
                    nbad = sum(1 for c in d.get("checks", []) if c["status"] != "ok")
                    lines.append(f"- {d['space']} ({d.get('likes', 0)} likes): {d['grade']}, {nbad} check(s) to fix")
            lines.append("Say which one to fix (e.g. \"diagnose 00000tt/xxx\") for details + add-ons.")
            return "\n".join(lines)
        if name == "state":
            return json.dumps(state_get(), default=str)
    except Exception as e:
        return f"Error running {name}: {e}"
    return f"Unknown tool: {name}"


def _local_fallback(messages: list[dict]) -> str:
    last = messages[-1].get("content", "") if messages else ""
    low = last.lower()
    if any(k in low for k in ("diagnose", "broken", "unreliable", "what's wrong", "whats wrong", "fix my")):
        if "mine" in low or "my space" in low or "my spaces" in low or low.strip() in ("diagnose", "diagnose mine"):
            res = _diagnose_mine()
            if res.get("error"):
                return res["error"]
            spaces = res.get("spaces") or []
            if not spaces:
                return "You have no published Spaces yet."
            lines = [f"I graded {res['count']} of your Spaces against the gold-standard rubric, worst first:"]
            for d in spaces[:6]:
                if d.get("error"):
                    continue
                nbad = sum(1 for c in d.get("checks", []) if c["status"] != "ok")
                lines.append(f"- {d['space']} ({d.get('likes', 0)} likes): {d['grade']}, {nbad} issue(s)")
            lines.append('Tell me one, e.g. "diagnose 00000tt/xxx", for the full checklist and add-ons.')
            return "\n".join(lines)
        m = re.search(r"(\w[\w./-]*/(\w[\w-]*))", last)
        space = m.group(1) if m else None
        if not space:
            return "Which Space should I diagnose? Give me an id like 00000tt/some-space (or say \"diagnose mine\")."
        d = _diagnose_space(space)
        if d.get("error"):
            return d["error"]
        lines = [f"{space} ({d['likes']} likes): GRADE = {d['grade']}"]
        for c in d["checks"]:
            mark = {"ok": "[ok]", "warn": "[~]", "missing": "[!!]"}.get(c["status"], "?")
            lines.append(f"{mark} {c['id']}: {c['evidence']}")
        bad = [c for c in d["checks"] if c["status"] != "ok"]
        if bad:
            lines.append("Fix order: " + ", ".join(c["id"] for c in bad[:5]) + ".")
        if d["addons"]:
            lines.append("Add-ons worth adding: " + "; ".join(d["addons"]) + ".")
        return "\n".join(lines)
    if any(k in low for k in ("compare", "similar", "find spaces", "other model", "workflow")):
        if "compare" in low or "app.py" in low:
            m = re.search(r"(\w[\w./-]*/(\w[\w-]*))", last)
            space = m.group(1) if m else None
            if not space:
                return "Tell me which Space to compare (id like user/space) and I'll pull its app.py, extract the workflow, and diff it against your config."
            info, path, text = _fetch_space(space)
            if not text:
                return f"Couldn't read app.py from {space} (private or no python file)."
            sig = _analyze_app(text)
            cmp = _compare_workflow(sig, State.config)
            out = f"{space} ({info.get('likes', 0)} likes): pipeline={sig['pipeline'] or 'unknown'}, base={sig['base_model'] or 'unknown'}, steps={sig.get('steps')}, guidance={sig.get('guidance')}, fps={sig.get('fps')}.\n"
            if cmp["deltas"]:
                out += "Notable differences:\n- " + "\n- ".join(cmp["deltas"]) + "\n\nWant me to refigure your config to match, or paste their app.py for you to look at?"
            else:
                out += "This workflow matches your config - nothing to change."
            return out
        res = hf_search_spaces("wan image to video")
        hits = [r for r in res if "error" not in r]
        if hits:
            top = sorted(hits, key=lambda r: r["likes"], reverse=True)[:5]
            lines = ["Similar working Spaces I found on the Hub:"]
            for r in top:
                lines.append(f"- {r['id']} · {r['likes']} likes · {r['sdk']}")
            lines.append("Tell me one to compare, e.g. \"compare user/space\", and I'll diff its app.py against your pipeline.")
            return "\n".join(lines)
        return "Search didn't turn up anything - try a more specific query."
    if "warehouse" in low or "list" in low or "asset" in low:
        a = warehouse_list()
        names = ", ".join(x["name"] for x in (a.get("assets") or [])[:6]) or "none"
        return f"Warehouse has {len(a.get('assets') or [])} LoRAs. First few: {names}."
    if "help" in low or "what can" in low:
        return ("I can: list warehouse assets, set pipeline config (base model, steps, "
                "guidance, LoRAs, space name), label example image+prompt pairs, publish "
                "a private Space, and flip visibility on published Spaces.")
    if "publish" in low:
        return "Publishing the current config as a private Space - I'll confirm when it's live."
    return ("I'm the Workshop assistant. Say 'list warehouse' to see the LoRA library, "
            "or describe the pipeline you want built (e.g. 'use FaceDownAssUp HIGH with "
            "6 steps') and I'll configure and publish it for you.")


@app.post("/api/chat")
def chat(payload: ChatMsg):
    if not State.vault_token:
        raise HTTPException(status_code=403, detail="Security vault empty. Enter a token first.")
    sys_prompt = SYSTEM_PROMPT.format(owner=OWNER, warehouse=WAREHOUSE_REPO)
    msgs = [{"role": "system", "content": sys_prompt}] + payload.messages
    reply = _chat_llm(msgs)
    if not reply:
        reply = _local_fallback(payload.messages)
        return {"role": "assistant", "content": reply}

    tool_call = _find_json_tool(reply)
    if tool_call:
        result = _execute_tool(tool_call)
        idx = reply.rfind("{")
        clean = reply[:idx].rstrip() if idx >= 0 else reply
        if result:
            clean = f"{clean}\n\n{result}"
        return {"role": "assistant", "content": clean}

    return {"role": "assistant", "content": reply}


if os.path.exists("./dist"):
    app.mount("/", StaticFiles(directory="./dist", html=True), name="static")