"""
edit_api.py — headless throughput API for Qwen-Image-Edit Rapid-AIO (v23 Q8 GGUF)
running on top of a local ComfyUI server.

Flow per request:  image + prompt  ->  upload to ComfyUI  ->  inject into the
workflow graph  ->  queue  ->  poll until done  ->  return the edited PNG.

Run ComfyUI first (run_comfyui.sh), then this service (start_api.sh).
"""

import io
import os
import json
import time
import uuid
import random
import copy
import threading
import csv

try:
    from . import database
    from . import embeddings
    from . import naming
except ImportError:
    import database
    import embeddings
    import naming

import requests
from PIL import Image
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import shutil
import re

# --- config -----------------------------------------------------------------
CONFIG_PATH  = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
WD_MODEL     = os.environ.get("WD_MODEL", "SmilingWolf/wd-vit-tagger-v3")
COMFY = os.environ.get("COMFY_URL", "http://127.0.0.1:8188").rstrip("/")
WORKFLOW_PATH = os.environ.get(
    "WORKFLOW_PATH",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow_qwen_edit.json"),
)
# Default target pixel area for the output latent. The MI50 is not fast, so we
# cap at ~1MP by default; raise via MAX_AREA env if you want bigger output.
MAX_AREA = int(os.environ.get("MAX_AREA", str(1024 * 1024)))
GEN_TIMEOUT = int(os.environ.get("GEN_TIMEOUT", "600"))  # seconds per request

# Node ids in workflow_qwen_edit.json (kept stable on purpose).
NODE_LOADIMAGE = "4"
NODE_POSITIVE = "5"
NODE_NEGATIVE = "6"
NODE_LATENT = "7"
NODE_KSAMPLER = "8"
NODE_SAVE = "10"

MAX_SEED = 2**32 - 1

VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.webm', '.mkv')

# Poses where the source image should be rotated 180° before pipeline for better results
ROTATE_180_POSES = {"the dragon", "dragon", "the draak", "draak"}

# WD tagger tags that indicate the subject is wearing clothes
CLOTHING_TAGS = {
    "dress", "skirt", "shirt", "top", "pants", "jeans", "jacket", "coat",
    "swimsuit", "bikini", "shorts", "uniform", "hoodie", "sweater", "blouse",
    "leggings", "stockings", "tights", "lingerie", "miniskirt", "pleated_skirt",
    "school_uniform", "maid_dress", "bodysuit", "sailor_uniform", "leotard",
    "corset", "kimono", "yukata", "cheongsam", "t-shirt", "tank_top", "crop_top",
    "tube_top", "halter_top", "negligee", "nightgown", "pajamas", "trench_coat",
    "vest", "bra", "underwear", "panties", "thong", "g-string", "bikini_top",
    "bikini_bottom", "one-piece_swimsuit", "sports_bra", "gym_clothes",
}

UNDRESS_PROMPT = "completely nude, bare skin, no clothing, naked body, natural skin texture"

with open(WORKFLOW_PATH, "r", encoding="utf-8") as f:
    BASE_WORKFLOW = json.load(f)

app = FastAPI(title="Qwen-Image-Edit Rapid-AIO API", version="1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "DELETE"],
    allow_headers=["*"],
)

def _sync_car_html():
    src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
    if not os.path.exists(src):
        return
    try:
        dest = os.path.join(_load_output_dir(), "car.html")
        shutil.copy2(src, dest)
        print(f"[car.html] synced → {dest}")
    except Exception as e:
        print(f"[car.html] sync warning: {e}")


def _watch_car_html():
    src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
    last_mtime = os.path.getmtime(src) if os.path.exists(src) else 0
    while True:
        time.sleep(1)
        try:
            mtime = os.path.getmtime(src)
            if mtime != last_mtime:
                last_mtime = mtime
                dest = os.path.join(_load_output_dir(), "car.html")
                shutil.copy2(src, dest)
                print(f"[car.html] change detected → synced to {dest}")
        except Exception:
            pass


def _load_wireframe_dir() -> str:
    with open(CONFIG_PATH, "r") as f:
        conf = json.load(f)
    return conf.get("wireframe_dir", "/mnt/zim/tour-comfy/wireframe")


def _load_faceswap_model_path() -> str:
    with open(CONFIG_PATH, "r") as f:
        conf = json.load(f)
    return os.path.expanduser(conf.get("faceswap_model", "~/.insightface/models/inswapper_128.onnx"))


@app.on_event("startup")
def on_startup():
    try:
        database.migrate_schema()
    except Exception as e:
        print(f"DB migration warning: {e}")
    _sync_car_html()
    threading.Thread(target=_watch_car_html, daemon=True).start()
    # Mount wireframe static dir for browser video preview
    try:
        wf_dir = _load_wireframe_dir()
        if os.path.isdir(wf_dir):
            app.mount("/wireframe", StaticFiles(directory=wf_dir), name="wireframe")
            print(f"[wireframe] mounted {wf_dir} → /wireframe")
    except Exception as e:
        print(f"[wireframe] mount warning: {e}")
    # Mount output dir so images can be served via HTTP (/output/filename.png)
    try:
        out_dir = _load_output_dir()
        if os.path.isdir(out_dir):
            app.mount("/output", StaticFiles(directory=out_dir), name="output")
            print(f"[output] mounted {out_dir} → /output")
    except Exception as e:
        print(f"[output] mount warning: {e}")
    # Write initial static data files (synchronous — ensures files exist before first request)
    _write_all_static()


# --- helpers ----------------------------------------------------------------
def _round16(x: int) -> int:
    return max(16, int(round(x / 16.0)) * 16)


def _target_size(w: int, h: int, max_area: int) -> tuple[int, int]:
    """Scale (w, h) to ~max_area preserving aspect, rounded to /16."""
    scale = (max_area / float(w * h)) ** 0.5
    return _round16(w * scale), _round16(h * scale)


def _prep_image(pil: Image.Image, max_area: int) -> tuple[Image.Image, int, int]:
    """
    Prepare image for ComfyUI:
    1. If area > max_area, crop from bottom if height remains >= 256.
    2. Otherwise scale (up or down) to fit area while preserving aspect.
    3. Ensure dimensions are rounded to 16.
    """
    w, h = pil.width, pil.height
    if w * h > max_area:
        # Try to keep width and crop height from bottom
        rw = _round16(w)
        th = max_area // rw
        if th >= 256:
            rh = (th // 16) * 16
            if rh < 16: rh = 16
            
            # To avoid black bars from .crop((0,0,rw,rh)) when rw > w,
            # we crop to original w first, then resize to rw.
            pil = pil.crop((0, 0, w, min(h, (rh * w) // rw)))
            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
            return pil, rw, rh
        else:
            # Too wide to keep width and have decent height, scale both down
            rw, rh = _target_size(w, h, max_area)
            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
            return pil, rw, rh
    else:
        # Fits or is too small: scale UP to match the max_area budget
        # (Legacy behavior that gives better model performance)
        rw, rh = _target_size(w, h, max_area)
        if rw != w or rh != h:
            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
        return pil, rw, rh


def _comfy_upload(img_bytes: bytes, filename: str) -> str:
    """Upload an image to ComfyUI's input dir; return the stored name."""
    r = requests.post(
        f"{COMFY}/upload/image",
        files={"image": (filename, img_bytes, "image/png")},
        data={"overwrite": "true", "type": "input"},
        timeout=60,
    )
    r.raise_for_status()
    j = r.json()
    name = j["name"]
    sub = j.get("subfolder", "")
    return f"{sub}/{name}" if sub else name


def _comfy_queue(graph: dict, client_id: str) -> str:
    r = requests.post(
        f"{COMFY}/prompt",
        json={"prompt": graph, "client_id": client_id},
        timeout=60,
    )
    if r.status_code != 200:
        raise HTTPException(502, f"ComfyUI rejected workflow: {r.text}")
    return r.json()["prompt_id"]


def _comfy_wait(prompt_id: str, deadline: float) -> dict:
    """Poll /history until the prompt finishes; return its outputs dict."""
    while time.time() < deadline:
        r = requests.get(f"{COMFY}/history/{prompt_id}", timeout=30)
        if r.status_code == 200:
            hist = r.json()
            if prompt_id in hist:
                entry = hist[prompt_id]
                status = entry.get("status", {})
                if status.get("status_str") == "error":
                    raise HTTPException(500, f"ComfyUI execution error: {json.dumps(status)}")
                outputs = entry.get("outputs", {})
                if outputs:
                    return outputs
        time.sleep(0.5)
    raise HTTPException(504, f"Generation timed out after {GEN_TIMEOUT}s")


def _comfy_fetch_image(outputs: dict) -> bytes:
    node_out = outputs.get(NODE_SAVE) or next(
        (v for v in outputs.values() if "images" in v), None
    )
    if not node_out or not node_out.get("images"):
        raise HTTPException(500, "No output image produced")
    img = node_out["images"][0]
    r = requests.get(
        f"{COMFY}/view",
        params={
            "filename": img["filename"],
            "subfolder": img.get("subfolder", ""),
            "type": img.get("type", "output"),
        },
        timeout=60,
    )
    r.raise_for_status()
    return r.content


# --- WD tagger (lazy) -------------------------------------------------------

_tagger      = None   # (model, transform, labels) once loaded
_tagger_lock = threading.Lock()


def _load_tagger():
    global _tagger
    if _tagger is not None:
        return _tagger
    with _tagger_lock:
        if _tagger is not None:
            return _tagger
        import torch
        import timm
        from timm.data import create_transform, resolve_data_config
        import huggingface_hub

        model = timm.create_model(f"hf_hub:{WD_MODEL}", pretrained=True).eval()
        if torch.cuda.is_available():
            model = model.cuda()

        cfg       = resolve_data_config(model.pretrained_cfg, model=model)
        transform = create_transform(**cfg)

        lpath  = huggingface_hub.hf_hub_download(WD_MODEL, "selected_tags.csv")
        with open(lpath, newline="") as f:
            rows   = list(csv.DictReader(f))
        # category 0=general 4=character 9=rating
        labels = [(r["name"], int(r.get("category", 9))) for r in rows]

        _tagger = (model, transform, labels)
    return _tagger


def _run_tagger(pil_img: Image.Image, threshold: float = 0.35):
    import torch
    model, transform, labels = _load_tagger()
    tensor = transform(pil_img.convert("RGB")).unsqueeze(0)
    if torch.cuda.is_available():
        tensor = tensor.cuda()
    with torch.no_grad():
        scores = torch.sigmoid(model(tensor))[0].cpu().tolist()
    tags = [
        {"tag": name, "score": round(score, 3), "cat": cat}
        for (name, cat), score in zip(labels, scores)
        if score >= threshold
    ]
    tags.sort(key=lambda x: -x["score"])
    return tags


def _tags_to_name(tags: list, max_tags: int = 8) -> str:
    content = [t["tag"] for t in tags if t["cat"] in (0, 4)][:max_tags]
    return " ".join(content).replace("_", " ")


def _apply_transparency(png_bytes: bytes) -> bytes:
    """Use rembg to remove background and return PNG bytes with Alpha channel."""
    try:
        from rembg import remove
        import io
        from PIL import Image
        img = Image.open(io.BytesIO(png_bytes))
        # rembg works best on RGB
        if img.mode != "RGB":
            img = img.convert("RGB")
        out = remove(img)
        buf = io.BytesIO()
        out.save(buf, format="PNG")
        return buf.getvalue()
    except Exception as e:
        print(f"Error in transparency post-processing: {e}")
        return png_bytes


# --- faceswapper (insightface + inswapper_128) --------------------------------
# Setup: pip install insightface onnxruntime-gpu opencv-python-headless
# Download: place inswapper_128.onnx at ~/.insightface/models/inswapper_128.onnx
# Source: https://huggingface.co/deepinsight/inswapper

_faceswapper      = None
_faceswapper_lock = threading.Lock()

# Dedicated single-worker pool for face-crop extraction.  Running it here
# instead of via FastAPI BackgroundTasks keeps the heavy insightface inference
# (and its one-time model load) off the shared request threadpool, so quick
# endpoints like /order stay responsive right after a "set preferred" click.
# A single worker also serializes face jobs so a burst can't thrash the GPU.
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
_face_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="face")

_gfpgan      = None
_gfpgan_lock = threading.Lock()


def _load_faceswapper():
    global _faceswapper
    if _faceswapper is not None:
        return _faceswapper
    with _faceswapper_lock:
        if _faceswapper is not None:
            return _faceswapper
        try:
            import insightface
            from insightface.app import FaceAnalysis
        except ImportError:
            raise RuntimeError("insightface not installed. Run: pip install insightface onnxruntime-gpu")

        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        app = FaceAnalysis(name='buffalo_l', providers=providers)
        app.prepare(ctx_id=0, det_size=(640, 640))

        model_path = _load_faceswap_model_path()
        if not os.path.exists(model_path):
            # Try HuggingFace download as fallback
            try:
                import huggingface_hub
                model_path = huggingface_hub.hf_hub_download(
                    'deepinsight/inswapper', 'inswapper_128.onnx',
                    local_dir=os.path.dirname(model_path),
                )
                print(f"[faceswap] Downloaded inswapper_128.onnx to {model_path}")
            except Exception as de:
                raise RuntimeError(
                    f"inswapper_128.onnx not found at {model_path}. "
                    f"Download from https://huggingface.co/deepinsight/inswapper and place it there. "
                    f"Download error: {de}"
                )

        swapper = insightface.model_zoo.get_model(model_path, providers=providers)
        _faceswapper = (app, swapper)
        print(f"[faceswap] loaded insightface buffalo_l + inswapper_128 from {model_path}")
    return _faceswapper


def _load_gfpgan():
    """Lazy-load GFPGAN face restorer. Returns restorer or False if unavailable."""
    global _gfpgan
    if _gfpgan is not None:
        return _gfpgan
    with _gfpgan_lock:
        if _gfpgan is not None:
            return _gfpgan
        try:
            from gfpgan import GFPGANer
            # Main GFPGAN model
            model_path = os.path.expanduser('~/.gfpgan/weights/GFPGANv1.4.pth')
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            if not os.path.exists(model_path):
                import urllib.request
                print('[gfpgan] Downloading GFPGANv1.4.pth (~333 MB)...')
                tmp = model_path + '.tmp'
                urllib.request.urlretrieve(
                    'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth',
                    tmp
                )
                os.rename(tmp, model_path)
            # GFPGANer hardcodes facexlib download path to CWD/gfpgan/weights/
            # → change CWD to ~ so models land at ~/gfpgan/weights/ (stable across runs)
            home = os.path.expanduser('~')
            os.makedirs(os.path.join(home, 'gfpgan', 'weights'), exist_ok=True)
            _orig_cwd = os.getcwd()
            os.chdir(home)
            try:
                restorer = GFPGANer(model_path=model_path, upscale=1, arch='clean',
                                    channel_multiplier=2, bg_upsampler=None)
            finally:
                os.chdir(_orig_cwd)
            _gfpgan = restorer
            print('[gfpgan] GFPGANv1.4 loaded')
        except Exception as e:
            print(f'[gfpgan] not available: {e}')
            _gfpgan = False
        return _gfpgan


def _make_video_poster(video_path: str) -> str | None:
    """Extract a poster JPG (sibling `<stem>.jpg`) so the gallery can show a
    thumbnail for a video via a plain <img> (file:// can't render <video> as a
    thumb). Returns the poster path on success, else None."""
    import subprocess
    poster_path = os.path.splitext(video_path)[0] + '.jpg'
    try:
        r = subprocess.run([
            'ffmpeg', '-y', '-ss', '1', '-i', video_path,
            '-frames:v', '1', '-q:v', '3', poster_path,
        ], capture_output=True, timeout=120)
        if r.returncode == 0 and os.path.exists(poster_path):
            return poster_path
        # -ss 1 can overshoot very short clips; retry from the first frame
        r = subprocess.run([
            'ffmpeg', '-y', '-i', video_path,
            '-frames:v', '1', '-q:v', '3', poster_path,
        ], capture_output=True, timeout=120)
        if r.returncode == 0 and os.path.exists(poster_path):
            return poster_path
    except Exception as pe:
        print(f'[poster] failed for {video_path}: {pe}')
    return None


def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance: bool = True, preview_scale: float = 1.0):
    """Frame-by-frame face swap: model face → every face in template video."""
    output_dir = _load_output_dir()
    wireframe_dir = _load_wireframe_dir()

    try:
        import cv2
        import numpy as np
        app, swapper = _load_faceswapper()
    except Exception as e:
        jobs[job_id]["status"] = "error"
        jobs[job_id]["error"] = str(e)
        return

    gfpgan_restorer = None
    if enhance:
        try:
            r = _load_gfpgan()
            if r is not False:
                gfpgan_restorer = r
        except Exception:
            pass

    try:
        # 1. Load source (model) face
        src_path = os.path.join(output_dir, model_filename)
        src_bgr = cv2.imread(src_path)
        if src_bgr is None:
            raise ValueError(f"Cannot read model image: {model_filename}")
        src_faces = app.get(src_bgr)
        if not src_faces:
            raise ValueError(f"No face detected in: {model_filename}")
        # Use the largest face as source
        src_face = max(src_faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))

        # 2. Open template video
        video_path = os.path.join(wireframe_dir, video_name)
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video: {video_name}")

        fps    = cap.get(cv2.CAP_PROP_FPS) or 25.0
        scale  = max(0.1, min(1.0, preview_scale))
        width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)  * scale)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale)
        # Ensure even dimensions for codec compatibility
        width  = width  if width  % 2 == 0 else width  - 1
        height = height if height % 2 == 0 else height - 1
        total  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        jobs[job_id]["total"] = max(total, 1)

        # 3. Write frame-swapped temp video
        ts        = time.strftime("%Y%m%d_%H%M%S")
        vid_stem  = os.path.splitext(video_name)[0]
        base_name = naming.get_base_name(model_filename)
        prev_tag  = f"_prev{int(scale*100)}" if scale < 1.0 else ""
        tmp_name  = f"{ts}_fs_tmp_{vid_stem}_{base_name}{prev_tag}.mp4"
        out_name  = f"{ts}_fs_{vid_stem}_{base_name}{prev_tag}.mp4"
        tmp_path  = os.path.join(output_dir, tmp_name)
        out_path  = os.path.join(output_dir, out_name)

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(tmp_path, fourcc, fps, (width, height))

        frame_idx = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            tgt_faces = app.get(frame)
            result = frame
            if tgt_faces:
                # Only swap the largest face — avoids false-positive detections
                # (reflections, background faces, face-like textures) causing ghost heads.
                # Ignore faces smaller than 40×40px (1600px²) as likely false positives.
                MIN_FACE_AREA = 1600
                valid = [f for f in tgt_faces
                         if (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]) >= MIN_FACE_AREA]
                if valid:
                    result = frame.copy()
                    best_face = max(valid, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
                    try:
                        result = swapper.get(result, best_face, src_face, paste_back=True)
                    except Exception:
                        pass
            if gfpgan_restorer is not None:
                try:
                    _, _, result = gfpgan_restorer.enhance(
                        result, has_aligned=False, only_center_face=False, paste_back=True
                    )
                except Exception:
                    pass
            if scale < 1.0:
                result = cv2.resize(result, (width, height), interpolation=cv2.INTER_AREA)
            writer.write(result)
            frame_idx += 1
            jobs[job_id]["done"] = frame_idx

        cap.release()
        writer.release()

        # 4. Remux with original audio via ffmpeg
        try:
            import subprocess
            r = subprocess.run([
                'ffmpeg', '-y',
                '-i', tmp_path,
                '-i', video_path,
                '-map', '0:v:0', '-map', '1:a?',
                '-c:v', 'libx264', '-preset', 'fast', '-crf', '18',
                '-c:a', 'aac', '-movflags', '+faststart',
                out_path,
            ], capture_output=True, timeout=600)
            if r.returncode == 0:
                os.remove(tmp_path)
            else:
                os.rename(tmp_path, out_path)
                print(f"[faceswap] ffmpeg failed ({r.returncode}), using raw mp4v output")
        except Exception as fe:
            os.rename(tmp_path, out_path)
            print(f"[faceswap] ffmpeg error: {fe}")

        # 5. Snapshot poster + register output in DB under same group as model
        _make_video_poster(out_path)
        person = database.get_person(model_filename)
        group_id = (person[1] if person and person[1] else naming.get_base_name(model_filename))
        database.upsert_person(
            out_name,
            filepath=out_path,
            group_id=group_id,
            content_type='video',
            faceswap_source_video=video_name,
            source_refs=json.dumps([model_filename]),
            sort_order=database.get_next_sort_order(group_id),
        )

        jobs[job_id]["status"] = "done"
        jobs[job_id]["output"] = out_name
        _invalidate_static()

    except Exception as e:
        print(f"[faceswap] error: {e}")
        jobs[job_id]["status"] = "error"
        jobs[job_id]["error"] = str(e)


def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
                         hair: bool = True, enhance: bool = True, preview_scale: float = 1.0):
    """High-quality faceswap via FaceFusion CLI (supports hair_swapper + ghost model)."""
    import subprocess as sp
    import sys
    import re as _re

    output_dir    = _load_output_dir()
    wireframe_dir = _load_wireframe_dir()

    with open(CONFIG_PATH, 'r') as f:
        conf = json.load(f)
    ff_dir  = os.path.expanduser(conf.get('facefusion_dir',  '~/facefusion'))
    ff_venv = os.path.expanduser(conf.get('facefusion_venv', '~/facefusion-venv'))

    ff_script = os.path.join(ff_dir, 'facefusion.py')
    ff_py     = os.path.join(ff_venv, 'bin', 'python')
    if not os.path.exists(ff_py):
        ff_py = sys.executable

    if not os.path.exists(ff_script):
        jobs[job_id]['status'] = 'error'
        jobs[job_id]['error']  = (
            f'FaceFusion not found at {ff_dir}. '
            'Run: bash tour-comfy/install_facefusion.sh'
        )
        return

    src_path   = os.path.join(output_dir,    model_filename)
    video_path = os.path.join(wireframe_dir, video_name)
    ts         = time.strftime('%Y%m%d_%H%M%S')
    vid_stem   = os.path.splitext(video_name)[0]
    base_name  = naming.get_base_name(model_filename)
    scale      = max(0.1, min(1.0, preview_scale))
    prev_tag   = f'_prev{int(scale*100)}' if scale < 1.0 else ''
    out_name   = f'{ts}_fs_{vid_stem}_{base_name}{prev_tag}.mp4'
    out_path   = os.path.join(output_dir, out_name)

    processors = ['face_swapper']
    # hair_swapper is not available in this FaceFusion version; use face_enhancer for quality
    if enhance:
        processors.append('face_enhancer')

    cmd = [
        ff_py, ff_script, 'headless-run',
        '--source-paths', src_path,
        '--target-path',  video_path,
        '--output-path', out_path,
        '--processors', *processors,
        '--execution-providers', 'cuda',
        # Limit to 2 threads: each thread owns a cuBLAS handle + workspace; more
        # threads exhausts VRAM when ComfyUI is running concurrently on the same GPU.
        '--execution-thread-count', '2',
        '--face-swapper-model', 'ghost_3_256',
        # The default yolo_face detector at score 0.5 misses the extreme-angle /
        # cropped close-up faces common in these POV template clips, so the swap
        # silently no-ops. scrfd at a lower score + multi-angle detection reliably
        # finds them; 'many' selector swaps every detected face per frame.
        '--face-detector-model', 'scrfd',
        '--face-detector-score', '0.3',
        '--face-detector-angles', '0', '90', '270',
        # 'one' swaps only the single best face per frame; 'many' caused ghost heads
        # by swapping false-positive detections (skin texture, reflections, etc.)
        '--face-selector-mode', 'one',
    ]
    if enhance:
        cmd += ['--face-enhancer-model', 'gfpgan_1.4']
    if scale < 1.0:
        # Determine native video width to compute preview width
        try:
            import cv2 as _cv2
            _cap = _cv2.VideoCapture(video_path)
            _native_w = int(_cap.get(_cv2.CAP_PROP_FRAME_WIDTH))
            _cap.release()
            _preview_w = max(2, int(_native_w * scale))
            if _preview_w % 2 != 0:
                _preview_w -= 1
            cmd += ['--output-video-width', str(_preview_w)]
        except Exception:
            pass

    jobs[job_id]['total'] = 100
    jobs[job_id]['done']  = 0

    # Ensure CUDA libs are on LD_LIBRARY_PATH for the subprocess (inherited from parent,
    # but also add nvidia package libs as fallback if running outside start_api.sh)
    import site as _site
    _sp_pkgs = next((p for p in _site.getsitepackages() if 'site-packages' in p), '')
    _nv_base = os.path.join(_sp_pkgs, 'nvidia')
    _extra_libs = ':'.join(
        os.path.join(_nv_base, pkg, 'lib')
        for pkg in ('cuda_runtime', 'cublas', 'cudnn', 'curand', 'cufft', 'cusolver', 'cusparse', 'nvjitlink', 'cuda_nvrtc')
        if os.path.isdir(os.path.join(_nv_base, pkg, 'lib'))
    )
    _env = os.environ.copy()
    if _extra_libs:
        _env['LD_LIBRARY_PATH'] = _extra_libs + (':' + _env['LD_LIBRARY_PATH'] if _env.get('LD_LIBRARY_PATH') else '')

    try:
        output_lines = []
        proc = sp.Popen(
            cmd, cwd=ff_dir, env=_env,
            stdout=sp.PIPE, stderr=sp.PIPE,
            text=True, errors='replace',
        )
        # FaceFusion writes tqdm progress to stderr; stdout carries other output.
        # Parse frame counts from both streams so the UI job counter updates.
        import threading as _thr
        def _parse_progress(line: str):
            m = _re.search(r'(\d+)\s*/\s*(\d+)', line)
            if m:
                done, total = int(m.group(1)), int(m.group(2))
                if total > 0:
                    jobs[job_id]['done']  = done
                    jobs[job_id]['total'] = total
        def _drain_stderr():
            for ln in proc.stderr:
                output_lines.append(ln.rstrip())
                print(f'[facefusion] {ln.rstrip()}')
                _parse_progress(ln)
        _thr.Thread(target=_drain_stderr, daemon=True).start()
        for line in proc.stdout:
            print(f'[facefusion] {line.rstrip()}')
            _parse_progress(line)
        proc.wait()

        if proc.returncode != 0:
            tail = '\n'.join(output_lines[-10:])
            raise RuntimeError(f'FaceFusion exited with code {proc.returncode}: {tail}')
        if not os.path.exists(out_path):
            raise RuntimeError('FaceFusion produced no output file')

        _make_video_poster(out_path)
        person   = database.get_person(model_filename)
        group_id = (person[1] if person and person[1] else naming.get_base_name(model_filename))
        database.upsert_person(
            out_name, filepath=out_path, group_id=group_id,
            content_type='video', faceswap_source_video=video_name,
            source_refs=json.dumps([model_filename]),
            sort_order=database.get_next_sort_order(group_id),
        )
        jobs[job_id]['status'] = 'done'
        jobs[job_id]['output'] = out_name
        _invalidate_static()

    except Exception as e:
        print(f'[faceswap-ff] error: {e}')
        jobs[job_id]['status'] = 'error'
        jobs[job_id]['error']  = str(e)


# --- pipeline helper ---------------------------------------------------------

def _load_poses():
    poses_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poses.md")
    if not os.path.exists(poses_path):
        return {}

    poses = {}
    current_pose = None
    current_beta = False
    current_desc = []

    with open(poses_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line.startswith("# "):
                if current_pose:
                    poses[current_pose] = {"text": " ".join(current_desc).strip(), "beta": current_beta}
                raw = line[2:].rstrip(":").strip()
                current_beta = bool(re.search(r'\(beta\)', raw, re.IGNORECASE))
                current_pose = re.sub(r'\s*\(beta\)\s*', '', raw, flags=re.IGNORECASE).strip()
                current_desc = []
            elif line and current_pose:
                current_desc.append(line)

    if current_pose:
        poses[current_pose] = {"text": " ".join(current_desc).strip(), "beta": current_beta}

    return poses


def _save_poses(poses: dict) -> None:
    """Rewrite poses.md from a {name: {text, beta}} dict, round-tripping _load_poses' format.

    Each pose is written as a ``# Name`` (or ``# Name (beta)``) header followed by its body.
    Body sentences separated by '. ' are written on their own lines to match the existing
    hand-authored style.
    """
    poses_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poses.md")
    blocks = []
    for name, entry in poses.items():
        name = str(name).strip()
        if not name:
            continue
        if isinstance(entry, dict):
            text = str(entry.get("text", "")).strip()
            beta = bool(entry.get("beta"))
        else:
            text = str(entry).strip()
            beta = False
        header = f"# {name}{' (beta)' if beta else ''}"
        # Split into readable lines on sentence boundaries without losing the period.
        body_lines = [s.strip() for s in re.split(r'(?<=\.)\s+', text) if s.strip()]
        blocks.append(header + "\n" + "\n".join(body_lines))
    with open(poses_path, "w", encoding="utf-8") as f:
        f.write("\n\n".join(blocks) + ("\n" if blocks else ""))


def _detect_has_background(pil: Image.Image) -> bool:
    """Return False when the image has significant alpha transparency (background removed)."""
    if pil.mode != 'RGBA':
        return True
    alpha = pil.split()[3]
    hist = alpha.histogram()
    transparent_px = sum(hist[:128])
    return transparent_px / (pil.width * pil.height) < 0.1


def _detect_has_clothing(tags: list) -> bool | None:
    """Return True if any tag from CLOTHING_TAGS appears above threshold, None if no tags."""
    if not tags:
        return None
    tag_names = {t["tag"] for t in tags}
    return bool(tag_names & CLOTHING_TAGS)


def _run_pipeline(
    pil: Image.Image,
    prompt: str,
    seed: int = -1,
    max_area: int = 0,
    steps: int = 4,
    cfg: float = 1.0,
    sampler_name: str = "euler_ancestral",
    scheduler: str = "beta",
    extra_images: list = None,  # additional PIL images wired to image2, image3
) -> bytes:
    area = max_area if max_area > 0 else MAX_AREA
    pil, w, h = _prep_image(pil, area)
    buf = io.BytesIO()
    pil.save(buf, format="PNG")
    stored = _comfy_upload(buf.getvalue(), f"in_{uuid.uuid4().hex[:8]}.png")
    if seed is None or seed < 0:
        seed = random.randint(0, MAX_SEED)
    graph = copy.deepcopy(BASE_WORKFLOW)
    graph[NODE_LOADIMAGE]["inputs"]["image"] = stored
    graph[NODE_POSITIVE]["inputs"]["prompt"] = prompt

    # Inject extra reference images as image2 / image3 on the positive encoder
    if extra_images:
        for i, extra_pil in enumerate(extra_images[:2]):
            extra_buf = io.BytesIO()
            extra_pil.convert("RGB").save(extra_buf, format="PNG")
            extra_stored = _comfy_upload(extra_buf.getvalue(), f"in_{uuid.uuid4().hex[:8]}.png")
            node_id = str(11 + i)   # "11" → image2, "12" → image3
            img_key = f"image{i + 2}"
            graph[node_id] = {
                "class_type": "LoadImage",
                "inputs": {"image": extra_stored},
                "_meta": {"title": f"ref image {i + 2}"},
            }
            graph[NODE_POSITIVE]["inputs"][img_key] = [node_id, 0]

    # ── background-removal routing ────────────────────────────────────────────
    # Two configurable strategies (config.json key "bg_removal"):
    #
    #   "rembg" (default) — strip transparent keyword → Qwen renders a natural
    #       scene → rembg (U2Net) separates person from any complex background.
    #
    #   "sam2" — replace transparent keyword with "black background" → Qwen
    #       renders a solid black BG → SAM2 bbox segmentation on a black image
    #       works perfectly because the contrast is maximal.
    #
    # Either way, explicit "black background" in the prompt always routes to
    # SAM2 (the user already set up the ideal SAM2 input).
    # ─────────────────────────────────────────────────────────────────────────
    _TRANSPARENT_KWS = ["transparent background", "no background",
                        "remove background", "alpha channel"]
    _BLACK_BG_KWS    = ["black background"]

    with open(CONFIG_PATH) as _cf:
        _bg_conf = json.load(_cf)
    bg_method = _bg_conf.get("bg_removal", "rembg")   # "rembg" | "sam2"

    is_transparent = any(kw in prompt.lower() for kw in _TRANSPARENT_KWS)
    is_black_bg    = any(kw in prompt.lower() for kw in _BLACK_BG_KWS)
    post_process   = None  # "rembg" | "sam2"

    if is_transparent:
        if bg_method == "sam2":
            # Swap "transparent background" → "black background" so Qwen renders
            # a pure-black BG that SAM2 can segment with maximal contrast.
            cleaned = prompt
            for kw in _TRANSPARENT_KWS:
                cleaned = re.sub(re.escape(kw), "black background", cleaned, flags=re.IGNORECASE)
            # Collapse duplicates if multiple keywords matched
            cleaned = re.sub(r"(?i)(black background[\s,]*){2,}", "black background, ", cleaned)
            cleaned = re.sub(r",\s*,", ",", cleaned).strip(", ")
            graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
            graph[NODE_NEGATIVE]["inputs"]["prompt"] = (
                "real background, outdoor scene, indoor scene, gradient, "
                "colored background, watermark, deformed anatomy"
            )
            post_process = "sam2"
        else:
            # Strip the keyword so Qwen renders a natural scene; rembg handles
            # any background complexity reliably.
            cleaned = prompt
            for kw in _TRANSPARENT_KWS:
                cleaned = re.sub(re.escape(kw), "", cleaned, flags=re.IGNORECASE)
            cleaned = re.sub(r",\s*,", ",", cleaned)
            cleaned = re.sub(r",\s*$", "", cleaned.strip()).strip(", ")
            graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
            graph[NODE_NEGATIVE]["inputs"]["prompt"] = "deformed anatomy, watermark, logo"
            post_process = "rembg"

    elif is_black_bg:
        # Prompt already specifies a black background — ideal SAM2 input.
        # Route to SAM2 regardless of the configured bg_removal method.
        post_process = "sam2"

    graph[NODE_LATENT]["inputs"]["width"] = w
    graph[NODE_LATENT]["inputs"]["height"] = h
    ks = graph[NODE_KSAMPLER]["inputs"]
    ks.update(seed=seed, steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler)
    client_id = uuid.uuid4().hex
    prompt_id = _comfy_queue(graph, client_id)
    outputs = _comfy_wait(prompt_id, time.time() + GEN_TIMEOUT)
    png_bytes = _comfy_fetch_image(outputs)

    if post_process == "sam2":
        # Input has a black background (Qwen was told "black background").
        # Use threshold-derived bbox so SAM2 gets a person-specific hint
        # rather than the full frame — full-frame bbox inverts the mask on
        # black-bg images because the large dark region scores higher.
        png_bytes = _apply_transparency_black_bg(png_bytes)
    elif post_process == "rembg":
        png_bytes = _apply_transparency(png_bytes)

    return png_bytes


# --- batch state -------------------------------------------------------------

jobs: dict[str, dict] = {}


def _load_output_dir() -> str:
    with open(CONFIG_PATH, "r") as f:
        conf = json.load(f)
    d = conf["output_dir"]
    if not os.path.isabs(d):
        d = os.path.normpath(os.path.join(os.path.dirname(CONFIG_PATH), "..", d))
    return d


def _move_to_trash(filepath: str):
    if not filepath or not os.path.exists(filepath):
        return
    output_dir = _load_output_dir()
    trash_dir = os.path.join(output_dir, ".trash")
    os.makedirs(trash_dir, exist_ok=True)
    
    filename = os.path.basename(filepath)
    ts = time.strftime("%Y%m%d_%H%M%S")
    trash_path = os.path.join(trash_dir, f"{ts}_{filename}")
    
    try:
        shutil.move(filepath, trash_path)
    except Exception as e:
        print(f"Error moving {filepath} to trash: {e}")


# --- static data files -------------------------------------------------------

_static_write_lock = threading.Lock()
_invalidate_timer: "threading.Timer | None" = None
_invalidate_timer_lock = threading.Lock()


def _write_json(path: str, data) -> None:
    """Atomic JSON write: write to .tmp then os.replace (no partial reads)."""
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(data, f)
    os.replace(tmp, path)


def _write_all_static() -> None:
    """Regenerate <output_dir>/_data/{images,names,groups,group-names,videos}.json.

    Called on startup and after every mutation so the browser can\
    fetch pre-computed JSON instead of hitting the DB on each page load.
    Uses a lock so concurrent invalidations don't race each other.
    """
    with _static_write_lock:
        try:
            output_dir = _load_output_dir()
            data_dir = os.path.join(output_dir, "_data")
            os.makedirs(data_dir, exist_ok=True)

            # Single DB query reused for images / names / groups
            persons = database.list_persons(include_archived=True)

            # images.json — all images (frontend filters by .archived field)
            db_images = []
            for p in persons:
                fpath = os.path.join(output_dir, p[0])
                if not os.path.exists(fpath):
                    continue
                db_images.append({
                    "filename":              p[0],
                    "name":                  p[1],
                    "group_id":              p[2],
                    "clip_description":      p[3],
                    "prompt":                p[4],
                    "pose":                  p[5],
                    "sort_order":            p[6],
                    "group_name":            p[7],
                    "hidden":                bool(p[8]) if p[8] else False,
                    "has_background":        bool(p[9]) if p[9] is not None else True,
                    "source_refs":           p[10],
                    "has_clothing":          p[11],
                    "content_type":          p[12] or "image",
                    "faceswap_source_video": p[13],
                    "archived":              bool(p[14]) if p[14] else False,
                })
            try:
                db_images.sort(
                    key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])),
                    reverse=True,
                )
            except Exception:
                pass
            _write_json(os.path.join(data_dir, "images.json"), {"images": db_images})

            # names.json — {filename: name}
            _write_json(os.path.join(data_dir, "names.json"),
                        {p[0]: p[1] for p in persons if p[1]})

            # groups.json — {filename: group_id}
            _write_json(os.path.join(data_dir, "groups.json"),
                        {p[0]: p[2] for p in persons if p[2]})

            # group-names.json — {group_id: display_name}
            _write_json(os.path.join(data_dir, "group-names.json"),
                        database.get_all_group_names())

            # videos.json — wireframe template videos
            wireframe_dir = _load_wireframe_dir()
            videos = []
            if os.path.isdir(wireframe_dir):
                videos = [f for f in sorted(os.listdir(wireframe_dir))
                          if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')]
            _write_json(os.path.join(data_dir, "videos.json"),
                        {"videos": videos, "wireframe_dir": wireframe_dir})

            # config.json — current config snapshot
            try:
                with open(CONFIG_PATH, "r") as _f:
                    _write_json(os.path.join(data_dir, "config.json"), json.load(_f))
            except Exception:
                pass

        except Exception as e:
            print(f"[static] write_all error: {e}")


def _invalidate_static() -> None:
    """Coalesce rapid invalidation calls — restarts a 0.3 s debounce timer each time.
    At most one _write_all_static() runs per quiet window, preventing thread floods
    during batch jobs that call this after every single image."""
    global _invalidate_timer
    with _invalidate_timer_lock:
        if _invalidate_timer is not None:
            _invalidate_timer.cancel()
        t = threading.Timer(0.3, _write_all_static)
        t.daemon = True
        t.start()
        _invalidate_timer = t


# -----------------------------------------------------------------------------

def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
                  seed: int, max_area: int, group_id: str | None = None,
                  wireframe_ref: str | None = None, wireframe_time: float = 0.5):
    output_dir = _load_output_dir()
    for fname in filenames:
        actual_gid = None
        try:
            person = database.get_person(fname)
            # Prefer the source's existing DB group_id; fall back to the caller-supplied
            # group_id (which is the gallery gid, potentially stale) or the basename.
            if person and person[1]:
                actual_gid = person[1]
            else:
                actual_gid = group_id or naming.get_base_name(fname)
            database.upsert_person(fname, group_id=actual_gid)
        except Exception as e:
            print(f"Error determining/updating group for {fname}: {e}")
            actual_gid = group_id or naming.get_base_name(fname)

        fpath = os.path.join(output_dir, fname)
        if not os.path.exists(fpath):
            jobs[job_id]["failed"] += len(prompts)
            continue

        try:
            base_pil = Image.open(fpath).convert("RGB")

            # Extract wireframe pose reference frame once per filename
            pose_guide_pil = None
            if wireframe_ref:
                try:
                    wf_path = os.path.join(_load_wireframe_dir(), wireframe_ref)
                    cap = cv2.VideoCapture(wf_path)
                    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                    target_frame = max(0, min(total_frames - 1, int(total_frames * wireframe_time)))
                    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
                    ret, frame = cap.read()
                    cap.release()
                    if ret:
                        pose_guide_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                        print(f"[batch] using wireframe {wireframe_ref} frame {target_frame}/{total_frames}")
                except Exception as wf_err:
                    print(f"[batch] wireframe extract error: {wf_err}")

            for prompt, pose in zip(prompts, poses):
                if jobs[job_id].get("cancelled"):
                    return
                try:
                    pil = base_pil
                    # Rotate 180° for poses that work better upside-down
                    if pose and pose.lower().strip() in ROTATE_180_POSES:
                        pil = pil.rotate(180)

                    extra_imgs = [pose_guide_pil] if pose_guide_pil else None
                    png = _run_pipeline(pil, prompt, seed, max_area, extra_images=extra_imgs)
                    ts = time.strftime("%Y%m%d_%H%M%S")
                    clean_fname = naming.get_base_name(fname)
                    out_name = f"{ts}_{clean_fname}"
                    out_path = os.path.join(output_dir, out_name)

                    with open(out_path, "wb") as f:
                        f.write(png)

                    has_bg = True
                    try:
                        out_pil = Image.open(io.BytesIO(png))
                        has_bg = _detect_has_background(out_pil)
                    except Exception:
                        pass

                    try:
                        embedding = embeddings.generate_embedding(out_path)
                        next_order = database.get_next_sort_order(actual_gid)
                        database.upsert_person(
                            out_name, filepath=out_path, embedding=embedding,
                            group_id=actual_gid, prompt=prompt, pose=pose,
                            has_background=has_bg, sort_order=next_order,
                            source_refs=json.dumps([fname]),
                        )
                    except Exception as db_err:
                        print(f"Database error in batch worker: {db_err}")

                    jobs[job_id]["done"] += 1
                    # Regenerate static JSON so the frontend's polling picks up the
                    # new image immediately (progressive refresh, not just at the end).
                    _invalidate_static()
                except Exception as e:
                    print(f"Error in batch for {fname} with prompt '{prompt}': {e}")
                    jobs[job_id]["failed"] += 1
        except Exception as e:
            print(f"Error opening {fname}: {e}")
            jobs[job_id]["failed"] += len(prompts)

    jobs[job_id]["status"] = "done"
    _invalidate_static()


def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], poses: list,
                      seed: int, max_area: int):
    """Generate one output image per prompt using filenames[0] as primary and the rest as extra refs."""
    output_dir = _load_output_dir()

    pils = []
    for fname in filenames:
        fpath = os.path.join(output_dir, fname)
        if os.path.exists(fpath):
            pils.append((fname, Image.open(fpath).convert("RGB")))

    if not pils:
        jobs[job_id]["status"] = "done"
        return

    # Output group: reuse shared group if all sources belong to the same one, else new group
    source_groups = set()
    for fname, _ in pils:
        try:
            p = database.get_person(fname)
            if p and p[1]:
                source_groups.add(p[1])
        except Exception:
            pass

    if len(source_groups) == 1:
        output_gid = next(iter(source_groups))
    else:
        output_gid = f"cg_{uuid.uuid4().hex[:8]}"

    primary_fname, primary_pil = pils[0]
    extra_pils = [p for _, p in pils[1:]]

    for prompt, pose in zip(prompts, poses):
        try:
            work_pil = primary_pil
            if pose and pose.lower().strip() in ROTATE_180_POSES:
                work_pil = work_pil.rotate(180)

            png = _run_pipeline(work_pil, prompt, seed, max_area, extra_images=extra_pils)
            ts = time.strftime("%Y%m%d_%H%M%S")
            clean = naming.get_base_name(primary_fname)
            out_name = f"{ts}_mr_{clean}"
            out_path = os.path.join(output_dir, out_name)

            with open(out_path, "wb") as f:
                f.write(png)

            has_bg = True
            try:
                out_pil = Image.open(io.BytesIO(png))
                has_bg = _detect_has_background(out_pil)
            except Exception:
                pass

            try:
                embedding = embeddings.generate_embedding(out_path)
                next_order = database.get_next_sort_order(output_gid)
                database.upsert_person(out_name, filepath=out_path, embedding=embedding,
                                       group_id=output_gid, prompt=prompt, pose=pose,
                                       has_background=has_bg, sort_order=next_order,
                                       source_refs=json.dumps([f for f, _ in pils]))
            except Exception as db_err:
                print(f"DB error in multi-ref: {db_err}")

            jobs[job_id]["done"] += 1
            # Regenerate static JSON so the frontend's polling picks up the new
            # image immediately (progressive refresh, matching _batch_worker).
            _invalidate_static()
        except Exception as e:
            print(f"Error in multi-ref for prompt '{prompt}': {e}")
            jobs[job_id]["failed"] += 1

    jobs[job_id]["status"] = "done"
    _invalidate_static()


# --- routes -----------------------------------------------------------------

class ConfigUpdate(BaseModel):
    prompt: str | None = None
    seed: int | None = None


@app.get("/config")
def get_config():
    with open(CONFIG_PATH, "r") as f:
        return json.load(f)


@app.post("/config")
def update_config(update: ConfigUpdate):
    with open(CONFIG_PATH, "r") as f:
        conf = json.load(f)
    if update.prompt is not None:
        conf["prompt"] = update.prompt
    if update.seed is not None:
        conf["seed"] = update.seed
    with open(CONFIG_PATH, "w") as f:
        json.dump(conf, f, indent=2)
    _invalidate_static()
    return {"prompt": conf["prompt"], "seed": conf["seed"]}


class BatchRequest(BaseModel):
    filenames: list[str]
    prompt: str | list[str]
    seed: int = -1
    max_area: int = 0
    group_id: str | None = None
    poses: list[str | None] | None = None  # pose name per prompt (same index), or None; None entries = no pose
    wireframe_ref: str | None = None        # wireframe video name to use as pose guide (image2 slot)
    wireframe_time: float = 0.5             # normalized time (0–1) to extract the pose frame from


@app.post("/batch")
def start_batch(req: BatchRequest):
    prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
    poses = req.poses or [None] * len(prompts)
    # Pad poses list to match prompts length
    while len(poses) < len(prompts):
        poses.append(None)
    total_tasks = len(req.filenames) * len(prompts)

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {"status": "running", "total": total_tasks, "done": 0, "failed": 0, "cancelled": False}
    t = threading.Thread(
        target=_batch_worker,
        args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
        kwargs={"wireframe_ref": req.wireframe_ref, "wireframe_time": req.wireframe_time},
        daemon=True,
    )
    t.start()
    return {"job_id": job_id, "total": total_tasks}


@app.delete("/batch/{job_id}")
def cancel_batch(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Job not found")
    jobs[job_id]["cancelled"] = True
    jobs[job_id]["status"] = "cancelled"
    return {"status": "cancelled", "job_id": job_id}


class MultiRefRequest(BaseModel):
    filenames: list[str]          # 2–3 reference images; first is primary (image1)
    prompt: str | list[str]
    poses: list[str | None] | None = None
    seed: int = -1
    max_area: int = 0


@app.post("/multi-ref")
def start_multi_ref(req: MultiRefRequest):
    if len(req.filenames) < 2:
        raise HTTPException(400, "multi-ref requires at least 2 filenames")
    filenames = req.filenames[:3]  # cap at 3 (image1/2/3)
    prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
    poses = req.poses or [None] * len(prompts)
    while len(poses) < len(prompts):
        poses.append(None)

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {"status": "running", "total": len(prompts), "done": 0, "failed": 0}
    t = threading.Thread(
        target=_multi_ref_worker,
        args=(job_id, filenames, prompts, poses, req.seed, req.max_area),
        daemon=True,
    )
    t.start()
    return {"job_id": job_id, "total": len(prompts)}


@app.get("/poses")
def get_poses():
    return _load_poses()


class PoseRequest(BaseModel):
    name: str
    text: str = ""
    beta: bool = False
    old_name: str | None = None  # set to rename an existing pose


@app.post("/poses")
def save_pose(req: PoseRequest):
    """Create, update, or rename a pose in poses.md."""
    name = req.name.strip()
    if not name:
        raise HTTPException(400, "Pose name is required")
    poses = _load_poses()
    # Rename: drop the old key (and preserve ordering by rebuilding).
    if req.old_name and req.old_name != name:
        poses.pop(req.old_name, None)
    poses[name] = {"text": req.text.strip(), "beta": bool(req.beta)}
    _save_poses(poses)
    return {"status": "success", "poses": poses}


@app.delete("/poses/{name}")
def delete_pose(name: str):
    """Delete a pose from poses.md."""
    poses = _load_poses()
    if name not in poses:
        raise HTTPException(404, "Pose not found")
    poses.pop(name, None)
    _save_poses(poses)
    return {"status": "success", "poses": poses}


@app.get("/batch/{job_id}")
def get_batch(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Job not found")
    return jobs[job_id]


@app.get("/images")
def list_images(archived: bool = False):
    output_dir = _load_output_dir()
    all_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg') + VIDEO_EXTENSIONS
    try:
        try:
            persons = database.list_persons(include_archived=archived)
            # list_persons cols: filename, name, group_id, clip_description,
            #   prompt, pose, sort_order, group_name, hidden, has_background,
            #   source_refs, has_clothing, content_type, faceswap_source_video, archived
            db_images = []
            for p in persons:
                fpath = os.path.join(output_dir, p[0])
                if not os.path.exists(fpath):
                    continue
                db_images.append({
                    "filename":             p[0],
                    "name":                 p[1],
                    "group_id":             p[2],
                    "clip_description":     p[3],
                    "prompt":               p[4],
                    "pose":                 p[5],
                    "sort_order":           p[6],
                    "group_name":           p[7],
                    "hidden":               bool(p[8]) if p[8] else False,
                    "has_background":       bool(p[9]) if p[9] is not None else True,
                    "source_refs":          p[10],
                    "has_clothing":         p[11],
                    "content_type":         p[12] or "image",
                    "faceswap_source_video":p[13],
                    "archived":             bool(p[14]) if p[14] else False,
                })
            db_images.sort(
                key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])),
                reverse=True,
            )
            return {"images": db_images}
        except Exception as db_err:
            print(f"DB error in list_images: {db_err}")
            listing = os.listdir(output_dir)
            # video poster snapshots share a video sibling's stem — don't list them as items
            video_stems = {os.path.splitext(f)[0] for f in listing if f.lower().endswith(VIDEO_EXTENSIONS)}
            files = [
                f for f in listing
                if f.lower().endswith(all_extensions)
                and not (f.lower().endswith('.jpg') and os.path.splitext(f)[0] in video_stems)
            ]
            files.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
            return {"images": [{"filename": f} for f in files]}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.get("/videos")
def list_videos():
    """Return available wireframe template videos."""
    wireframe_dir = _load_wireframe_dir()
    if not os.path.isdir(wireframe_dir):
        return {"videos": []}
    videos = [
        f for f in sorted(os.listdir(wireframe_dir))
        if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')
    ]
    return {"videos": videos, "wireframe_dir": wireframe_dir}


@app.get("/wireframe/frame/{video_name}")
def wireframe_frame(video_name: str, t: float = 0.5):
    """Extract a single frame at normalized time t (0–1) from a wireframe video. Returns PNG."""
    import cv2
    wireframe_dir = _load_wireframe_dir()
    video_path = os.path.join(wireframe_dir, video_name)
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Video not found: {video_name}")
    try:
        cap = cv2.VideoCapture(video_path)
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        target = max(0, min(total - 1, int(total * max(0.0, min(1.0, t)))))
        cap.set(cv2.CAP_PROP_POS_FRAMES, target)
        ret, frame = cap.read()
        cap.release()
        if not ret:
            raise HTTPException(500, "Could not read frame")
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil = Image.fromarray(rgb)
        buf = io.BytesIO()
        pil.save(buf, format="PNG")
        buf.seek(0)
        from fastapi.responses import Response
        return Response(content=buf.getvalue(), media_type="image/png")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, f"Frame extraction error: {e}")


@app.get("/wireframe/duration/{video_name}")
def wireframe_duration(video_name: str):
    """Return duration (seconds) of a wireframe video via ffprobe."""
    import subprocess
    wireframe_dir = _load_wireframe_dir()
    video_path = os.path.join(wireframe_dir, video_name)
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Video not found: {video_name}")
    try:
        r = subprocess.run(
            ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
             '-of', 'json', video_path],
            capture_output=True, timeout=10,
        )
        info = json.loads(r.stdout)
        duration = float(info.get('format', {}).get('duration', 0))
    except Exception as e:
        raise HTTPException(500, f"ffprobe error: {e}")
    return {'video_name': video_name, 'duration': duration}


class TrimRequest(BaseModel):
    video_name: str
    start: float                  # seconds
    end: float                    # seconds
    output_name: str | None = None  # auto-generated if None


@app.post("/wireframe/trim")
def trim_wireframe(req: TrimRequest):
    """Trim a wireframe video to [start, end] seconds using ffmpeg stream copy."""
    import subprocess
    wireframe_dir = _load_wireframe_dir()
    video_path = os.path.join(wireframe_dir, req.video_name)
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Video not found: {req.video_name}")
    if req.start < 0 or req.end <= req.start:
        raise HTTPException(400, "Invalid start/end: end must be > start ≥ 0")

    stem = os.path.splitext(req.video_name)[0]
    if req.output_name:
        out_name = req.output_name if req.output_name.lower().endswith('.mp4') else req.output_name + '.mp4'
    else:
        out_name = f"{stem}_{int(req.start)}s-{int(req.end)}s.mp4"

    out_path = os.path.join(wireframe_dir, out_name)
    if os.path.exists(out_path):
        raise HTTPException(409, f"File already exists: {out_name}")

    r = subprocess.run(
        ['ffmpeg', '-y',
         '-ss', str(req.start), '-to', str(req.end),
         '-i', video_path,
         '-c', 'copy',
         out_path],
        capture_output=True, timeout=120,
    )
    if r.returncode != 0:
        raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:500]}")

    return {'output_name': out_name, 'start': req.start, 'end': req.end}


class VideoStitchRequest(BaseModel):
    filenames: list[str]        # images from output_dir, in desired frame order
    fps: float = 8.0
    loop: bool = True           # append reversed sequence for a ping-pong loop
    output_name: str | None = None


@app.post("/generate-video")
def generate_video(req: VideoStitchRequest):
    """Stitch a list of output images into a short looping MP4 using ffmpeg."""
    import subprocess, tempfile, textwrap
    if len(req.filenames) < 2:
        raise HTTPException(400, "Need at least 2 images")
    if not (0.1 <= req.fps <= 60):
        raise HTTPException(400, "fps must be between 0.1 and 60")

    output_dir = _load_output_dir()

    # Validate + collect paths
    paths = []
    for fn in req.filenames:
        p = os.path.join(output_dir, fn)
        if not os.path.exists(p):
            raise HTTPException(404, f"Image not found: {fn}")
        paths.append(p)

    # Ping-pong loop: forward + reversed (drop first/last to avoid duplicate frames at seam)
    if req.loop and len(paths) > 1:
        paths = paths + list(reversed(paths[1:-1])) if len(paths) > 2 else paths + list(reversed(paths))

    ts = time.strftime("%Y%m%d_%H%M%S")
    base = naming.get_base_name(req.filenames[0])
    stem = os.path.splitext(base)[0]
    out_name = req.output_name or f"{ts}_vid_{stem}.mp4"
    if not out_name.lower().endswith(".mp4"):
        out_name += ".mp4"
    out_path = os.path.join(output_dir, out_name)

    # Build ffmpeg concat list in a temp file
    duration = 1.0 / req.fps
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as tf:
        for p in paths:
            tf.write(f"file '{p}'\nduration {duration:.6f}\n")
        # ffmpeg concat needs the last entry without a duration
        tf.write(f"file '{paths[-1]}'\n")
        concat_path = tf.name

    try:
        r = subprocess.run(
            ["ffmpeg", "-y",
             "-f", "concat", "-safe", "0", "-i", concat_path,
             "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2",   # ensure even dimensions
             "-c:v", "libx264", "-pix_fmt", "yuv420p",
             "-movflags", "+faststart",
             out_path],
            capture_output=True, timeout=120,
        )
    finally:
        os.unlink(concat_path)

    if r.returncode != 0:
        raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:600]}")

    # Register in DB so it shows up in the gallery
    person = database.get_person(req.filenames[0])
    group_id = person[1] if person and person[1] else naming.get_base_name(req.filenames[0])
    try:
        next_order = database.get_next_sort_order(group_id)
        database.upsert_person(
            out_name, filepath=out_path,
            group_id=group_id,
            sort_order=next_order,
            content_type="video",
            source_refs=json.dumps(req.filenames),
        )
        _invalidate_static()
    except Exception as e:
        print(f"[generate-video] DB error: {e}")

    return {"output": out_name, "frames": len(paths), "fps": req.fps}


class FrameExtractRequest(BaseModel):
    video_name: str
    time: float = 0.0


@app.post("/wireframe/frame")
def wireframe_extract_frame(req: FrameExtractRequest):
    """Extract a single frame at a given timestamp and return it as base64 PNG."""
    import base64
    wireframe_dir = _load_wireframe_dir()
    video_path = os.path.join(wireframe_dir, req.video_name)
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Video not found: {req.video_name}")
    try:
        img = _extract_frame_at(video_path, req.time)
    except Exception as e:
        raise HTTPException(500, str(e))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return {"frame_b64": base64.b64encode(buf.getvalue()).decode()}


class FaceswapRequest(BaseModel):
    model_filename: str        # image from output_dir to use as face source
    video_name: str            # filename of template video in wireframe_dir
    enhance: bool = True       # GFPGAN face restoration after each frame swap
    hair: bool = False         # use FaceFusion with hair_swapper (requires FaceFusion install)
    preview_scale: float = 1.0 # 0.25–1.0; <1.0 produces a smaller preview video


@app.get("/faceswap/check")
def faceswap_check():
    """Report which enhancement backends are available."""
    gfpgan_ok = False
    try:
        import gfpgan  # noqa
        gfpgan_ok = True
    except ImportError:
        pass

    with open(CONFIG_PATH, 'r') as f:
        conf = json.load(f)
    ff_dir    = os.path.expanduser(conf.get('facefusion_dir', '~/facefusion'))
    ff_script = os.path.join(ff_dir, 'facefusion.py')

    return {'gfpgan': gfpgan_ok, 'facefusion': os.path.exists(ff_script)}


@app.post("/faceswap")
def start_faceswap(req: FaceswapRequest):
    output_dir    = _load_output_dir()
    wireframe_dir = _load_wireframe_dir()

    src_path   = os.path.join(output_dir,    req.model_filename)
    video_path = os.path.join(wireframe_dir, req.video_name)

    if not os.path.exists(src_path):
        raise HTTPException(404, f"Model image not found: {req.model_filename}")
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Template video not found: {req.video_name}")

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {
        "status": "running", "type": "faceswap",
        "total": 1, "done": 0, "failed": 0,
        "model": req.model_filename, "video": req.video_name,
    }

    if req.hair:
        t = threading.Thread(
            target=_faceswap_worker_ff,
            args=(job_id, req.model_filename, req.video_name),
            kwargs={'hair': True, 'enhance': req.enhance, 'preview_scale': req.preview_scale},
            daemon=True,
        )
    else:
        t = threading.Thread(
            target=_faceswap_worker,
            args=(job_id, req.model_filename, req.video_name),
            kwargs={'enhance': req.enhance, 'preview_scale': req.preview_scale},
            daemon=True,
        )
    t.start()
    return {"job_id": job_id, "model": req.model_filename, "video": req.video_name}


# --- tagging routes ----------------------------------------------------------

class TagRequest(BaseModel):
    filename: str
    threshold: float = 0.35
    max_tags: int = 8
    group_id: str | None = None


@app.post("/tag")
def tag_image(req: TagRequest):
    output_dir = _load_output_dir()
    fpath = os.path.join(output_dir, req.filename)
    if not os.path.exists(fpath):
        raise HTTPException(404, "File not found in output dir")
    try:
        pil  = Image.open(fpath)
        tags = _run_tagger(pil, req.threshold)
        clip_desc = _tags_to_name(tags, req.max_tags)
        has_clothing = _detect_has_clothing(tags)

        # Only assign a new name if the image doesn't already have one
        existing = database.get_person(req.filename)
        auto_name = (existing[0] if existing and existing[0] else None) or naming.generate_associative_name(tags)

        # Save to DB
        try:
            embedding = embeddings.generate_embedding(fpath)
            database.upsert_person(
                req.filename, filepath=fpath, name=auto_name,
                clip_description=clip_desc, tags=tags, embedding=embedding,
                group_id=req.group_id, has_clothing=has_clothing,
            )
        except Exception as db_err:
            print(f"Database error during tag: {db_err}")

        return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30], "has_clothing": has_clothing}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.get("/names")
def get_names():
    try:
        persons = database.list_persons()
        return {p[0]: p[1] for p in persons if p[1]}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/names/{filename}")
def set_name(filename: str, body: dict):
    name = body.get("name", "")
    try:
        database.upsert_person(filename, name=name)
    except Exception as db_err:
        print(f"Database error in set_name: {db_err}")

    _invalidate_static()
    return {"filename": filename, "name": name}


# --- group routes ------------------------------------------------------------

@app.get("/groups")
def get_groups():
    try:
        persons = database.list_persons()
        return {p[0]: p[2] for p in persons if p[2]}
    except Exception as e:
        raise HTTPException(500, str(e))


class MergeRequest(BaseModel):
    filenames: list[str]
    group_id: str | None = None


@app.post("/groups/merge")
def merge_groups(req: MergeRequest):
    gid = req.group_id or f"cg_{uuid.uuid4().hex[:8]}"
    for fname in req.filenames:
        try:
            database.upsert_person(fname, group_id=gid)
        except Exception as db_err:
            print(f"Database error in merge: {db_err}")

    # Write synchronously: the frontend reloads images.json immediately after this
    # returns, so an async rebuild would race and show the pre-merge grouping.
    _write_all_static()
    return {"group_id": gid, "files": req.filenames}


class ExtractRequest(BaseModel):
    filename: str


@app.post("/groups/extract")
def extract_from_group(req: ExtractRequest):
    gid = f"solo:{req.filename}"
    try:
        database.upsert_person(req.filename, group_id=gid)
    except Exception as db_err:
        print(f"Database error in extract: {db_err}")

    _write_all_static()
    return {"filename": req.filename}


@app.get("/group-names")
def get_group_names():
    try:
        return database.get_all_group_names()
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/group-names/{group_id}")
def set_group_name(group_id: str, body: dict):
    name = body.get("name", "").strip()
    try:
        database.set_group_name(group_id, name or None)
    except Exception as e:
        raise HTTPException(500, str(e))
    _invalidate_static()
    return {"group_id": group_id, "name": name}


@app.get("/groups/{group_id}/order")
def get_group_order(group_id: str):
    try:
        rows = database.get_group_order(group_id)
        return {"group_id": group_id, "filenames": [r[0] for r in rows]}
    except Exception as e:
        raise HTTPException(500, str(e))


class GroupOrderRequest(BaseModel):
    filenames: list[str]


@app.post("/groups/{group_id}/order")
def set_group_order(group_id: str, req: GroupOrderRequest):
    try:
        database.set_group_order(group_id, req.filenames)
    except Exception as e:
        raise HTTPException(500, str(e))
    _invalidate_static()
    return {"group_id": group_id, "filenames": req.filenames}


@app.get("/similar/{filename}")
def get_similar(filename: str, limit: int = 10):
    person = database.get_person(filename)
    if not person or person[3] is None:
        raise HTTPException(404, "Image or embedding not found")
    
    embedding = person[3]
    results = database.search_similar(embedding, limit=limit)
    
    similar = []
    for r in results:
        # Avoid returning the same image as the most similar
        if r[0] == filename:
            continue
        similar.append({
            "filename": r[0],
            "name": r[1],
            "group_id": r[2],
            "clip_description": r[3],
            "distance": float(r[4])
        })
    
    return {"filename": filename, "similar": similar}


@app.post("/db/cleanup")
def db_cleanup():
    """Delete DB records for files that no longer exist on disk."""
    output_dir = _load_output_dir()
    persons = database.list_persons()
    removed = []
    for p in persons:
        fpath = os.path.join(output_dir, p[0])
        if not os.path.exists(fpath):
            database.delete_person(p[0])
            removed.append(p[0])
    if removed:
        _invalidate_static()
    return {"removed": len(removed), "filenames": removed}


@app.get("/health")
def health():
    try:
        requests.get(f"{COMFY}/system_stats", timeout=5).raise_for_status()
        return {"status": "ok", "comfy": COMFY}
    except Exception as e:
        raise HTTPException(503, f"ComfyUI unreachable at {COMFY}: {e}")


def _crop_to_bbox(pil_img: Image.Image, margin: int = 20, top_margin: int = 20, headroom: float = 0.05) -> Image.Image:
    if pil_img.mode != 'RGBA':
        return pil_img
    
    alpha = pil_img.split()[-1]
    bbox = alpha.getbbox()
    if not bbox:
        return pil_img
    
    left, upper, right, lower = bbox
    left = max(0, left - margin)
    upper = max(0, upper - top_margin)
    right = min(pil_img.width, right + margin)
    lower = min(pil_img.height, lower + margin)
    
    cropped = pil_img.crop((left, upper, right, lower))
    if headroom > 0:
        h_px = int(cropped.height * headroom)
        if h_px > 0:
            new_img = Image.new("RGBA", (cropped.width, cropped.height + h_px), (0, 0, 0, 0))
            new_img.paste(cropped, (0, h_px))
            return new_img
    return cropped


def _extract_face_bg(filename: str, fpath: str):
    """Background task: detect largest face, crop with padding, save as {group_id}_face.png."""
    try:
        import cv2
        app_fa, _ = _load_faceswapper()
        bgr = cv2.imread(fpath)
        if bgr is None:
            print(f"[extract-face] cannot read {fpath}")
            return
        faces = app_fa.get(bgr)
        if not faces:
            print(f"[extract-face] no face detected in {filename}")
            return
        face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
        x1, y1, x2, y2 = [int(v) for v in face.bbox]
        h, w = bgr.shape[:2]
        pad = int((y2 - y1) * 0.5)
        x1 = max(0, x1 - pad)
        y1 = max(0, y1 - pad * 2)   # extra headroom above face
        x2 = min(w, x2 + pad)
        y2 = min(h, y2 + int(pad * 0.3))
        pil = Image.open(fpath).convert("RGBA")
        cropped = pil.crop((x1, y1, x2, y2))
        person = database.get_person(filename)
        group_id = person[1] if person else None
        gid_tag = (group_id or "face").replace("/", "_")
        face_fname = f"{gid_tag}_face.png"
        face_path = os.path.join(os.path.dirname(fpath), face_fname)
        cropped.save(face_path)
        face_embed = face.normed_embedding.tolist() if hasattr(face, 'normed_embedding') and face.normed_embedding is not None else None
        database.upsert_person(face_fname, filepath=face_path, group_id=group_id,
                               name=person[0] if person else None,
                               source_refs=json.dumps([filename]),
                               face_embedding=face_embed)
        print(f"[extract-face] saved {face_fname}" + (" + face embedding" if face_embed else ""))
    except Exception as e:
        print(f"[extract-face] error for {filename}: {e}")


def _process_upload(file_path: str, filename: str, prompts: list[str], name: str | None = None, group_id: str | None = None, poses: list[str] | None = None):
    output_dir = _load_output_dir()
    try:
        pil = Image.open(file_path)
        
        # 1. CLIP tag the source
        tags = _run_tagger(pil.convert("RGB"))
        clip_desc = _tags_to_name(tags)
        has_clothing = _detect_has_clothing(tags)
        auto_name = name or naming.generate_associative_name(tags)

        # 2. Embedding for source
        embedding = embeddings.generate_embedding(file_path)

        # 3. Register source in DB — sort_order=0 makes it the preferred base image
        database.upsert_person(
            filename, filepath=file_path, name=auto_name,
            clip_description=clip_desc, tags=tags, embedding=embedding,
            group_id=group_id, sort_order=0, has_clothing=has_clothing,
        )
        # Surface the new group with its base image right away — the pose/base-prompt
        # generation below can take a while, and the user shouldn't wait for it to
        # see the group land on the gallery.
        _invalidate_static()

        # 4. Crop if needed
        cropped_pil = _crop_to_bbox(pil)

        # 5. Run prompts
        for i, prompt in enumerate(prompts):
            try:
                png = _run_pipeline(cropped_pil.convert("RGB"), prompt)

                ts = time.strftime("%Y%m%d_%H%M%S")
                out_name = f"{ts}_{i}_{filename}"
                if not out_name.lower().endswith(".png"):
                    out_name += ".png"
                out_path = os.path.join(output_dir, out_name)

                with open(out_path, "wb") as f:
                    f.write(png)

                out_embedding = embeddings.generate_embedding(out_path)
                next_order = database.get_next_sort_order(group_id)
                # Persist the prompt (and pose name, when this output came from a named pose)
                # so the generation parameters survive in the DB / images.json.
                out_pose = poses[i] if (poses and i < len(poses)) else None
                database.upsert_person(
                    out_name, filepath=out_path, name=auto_name,
                    clip_description=clip_desc, embedding=out_embedding,
                    group_id=group_id, sort_order=next_order,
                    prompt=prompt, pose=out_pose,
                )
            except Exception as e:
                print(f"Error processing prompt '{prompt}' for {filename}: {e}")
                
    except Exception as e:
        print(f"Error in _process_upload for {filename}: {e}")
    finally:
        _invalidate_static()


@app.post("/upload")
def upload_image(
    background_tasks: BackgroundTasks,
    image: UploadFile = File(...),
    prompts: str = Form(""),
    name: str = Form(None),
    group_id: str = Form(None),      # optional: add to existing group
    skip_poses: bool = Form(False),  # optional: skip base_prompts generation
):
    # Load config to get output_dir (we use output_dir for UI uploads to avoid watcher conflict)
    with open(CONFIG_PATH, "r") as f:
        conf = json.load(f)
    output_dir = _load_output_dir()
    os.makedirs(output_dir, exist_ok=True)

    ts = time.strftime("%Y%m%d_%H%M%S")
    safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename or "paste")
    # Ensure extension
    if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
        safe_filename += ".png"

    filename = f"{ts}_{safe_filename}"
    file_path = os.path.join(output_dir, filename)

    with open(file_path, "wb") as f:
        shutil.copyfileobj(image.file, f)

    # Fast path: add to existing group without pose generation
    if group_id and skip_poses:
        sort_order = database.get_next_sort_order(group_id)
        database.upsert_person(filename, filepath=file_path, group_id=group_id,
                               sort_order=sort_order)
        _invalidate_static()
        return {"status": "added", "filename": filename, "group_id": group_id}

    prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]

    # Add base-set prompts if defined in config
    base_prompts = conf.get("base_prompts", [])
    if isinstance(base_prompts, list):
        prompt_list.extend(base_prompts)

    if not prompt_list:
        # Use default prompt from config
        prompt_list = [conf.get("prompt", "high quality, masterpiece")]

    effective_gid = group_id or f"up_{uuid.uuid4().hex[:8]}"
    background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, effective_gid)

    return {"status": "processing", "filename": filename, "group_id": effective_gid, "prompts": prompt_list}


@app.post("/edit")
async def edit(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    seed: int = Form(-1),
    steps: int = Form(4),
    cfg: float = Form(1.0),
    sampler_name: str = Form("euler_ancestral"),
    scheduler: str = Form("beta"),
    max_area: int = Form(0),
):
    raw = await image.read()
    try:
        pil = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception as e:
        raise HTTPException(400, f"Invalid image: {e}")

    png = _run_pipeline(pil, prompt, seed, max_area, steps, cfg, sampler_name, scheduler)
    return Response(content=png, media_type="image/png")


@app.post("/images/{filename}/hidden")
def set_image_hidden(filename: str, body: dict):
    hidden = bool(body.get("hidden", False))
    try:
        database.set_hidden(filename, hidden)
    except Exception as e:
        raise HTTPException(500, str(e))
    _invalidate_static()
    return {"filename": filename, "hidden": hidden}


@app.post("/images/{filename}/archive")
def archive_image(filename: str):
    try:
        database.set_archived(filename, True)
    except Exception as e:
        raise HTTPException(500, str(e))
    _invalidate_static()
    return {"filename": filename, "archived": True}


@app.post("/images/{filename}/unarchive")
def unarchive_image(filename: str):
    try:
        database.set_archived(filename, False)
    except Exception as e:
        raise HTTPException(500, str(e))
    _invalidate_static()
    return {"filename": filename, "archived": False}


@app.post("/images/{filename}/set-preferred")
def set_image_preferred(filename: str):
    """Make this image sort_order=0 within its group, shifting others to 1,2,..."""
    person = database.get_person(filename)
    if not person:
        raise HTTPException(404, "Image not found")
    group_id = person[1]
    if not group_id:
        raise HTTPException(400, "Image has no group assigned")
    rows = database.get_group_order(group_id)
    others = [r[0] for r in rows if r[0] != filename]
    database.set_group_order(group_id, [filename] + others)
    _invalidate_static()
    fpath = os.path.join(_load_output_dir(), filename)
    if os.path.exists(fpath):
        _face_executor.submit(_extract_face_bg, filename, fpath)
    return {"filename": filename, "group_id": group_id}


@app.post("/images/{filename}/extract-face")
def extract_face_endpoint(filename: str):
    """Detect and crop the largest face from image; saves as {group_id}_face.png."""
    output_dir = _load_output_dir()
    fpath = os.path.join(output_dir, filename)
    if not os.path.exists(fpath):
        raise HTTPException(404, "not found")
    _face_executor.submit(_extract_face_bg, filename, fpath)
    return {"status": "queued", "filename": filename}


class FaceSimilarRequest(BaseModel):
    group_id: str
    limit: int = 12


@app.post("/faces/similar")
def face_similar(req: FaceSimilarRequest):
    """Find groups with visually similar faces using insightface embeddings.

    Looks up the face embedding stored for {group_id}_face.png and returns
    the top-N closest matches from other groups.
    """
    face_fname = f"{req.group_id.replace('/', '_')}_face.png"
    embedding = database.get_face_embedding(face_fname)
    if embedding is None:
        raise HTTPException(404, "No face embedding found for this group — set a preferred image first")

    rows = database.search_similar_face(embedding, limit=req.limit, exclude_group_id=req.group_id)
    # Each row is (filename, group_id, distance). Return the group thumbnail filename
    # (the _face.png itself) so the frontend can render it directly.
    results = [
        {"filename": r[0], "group_id": r[1], "distance": round(float(r[2]), 4)}
        for r in rows
    ]
    return {"similar": results}


_face_index_status: dict = {"running": False, "done": 0, "total": 0, "indexed": 0}


def _face_index_worker():
    """Backfill face embeddings for all *_face.png files that lack one."""
    global _face_index_status
    output_dir = _load_output_dir()
    face_files = [f for f in os.listdir(output_dir) if f.endswith("_face.png")]
    _face_index_status.update({"running": True, "done": 0, "total": len(face_files), "indexed": 0})
    try:
        import cv2
        app_fa, _ = _load_faceswapper()
    except Exception as e:
        print(f"[face-index] failed to load insightface: {e}")
        _face_index_status["running"] = False
        return
    indexed = 0
    for i, fname in enumerate(face_files):
        existing = database.get_face_embedding(fname)
        if existing is not None:
            _face_index_status["done"] = i + 1
            continue
        fpath = os.path.join(output_dir, fname)
        try:
            bgr = cv2.imread(fpath)
            if bgr is None:
                continue
            faces = app_fa.get(bgr)
            if not faces:
                continue
            face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
            if not hasattr(face, 'normed_embedding') or face.normed_embedding is None:
                continue
            database.upsert_person(fname, face_embedding=face.normed_embedding.tolist())
            indexed += 1
        except Exception as e:
            print(f"[face-index] {fname}: {e}")
        _face_index_status.update({"done": i + 1, "indexed": indexed})
    _face_index_status["running"] = False
    print(f"[face-index] done: {indexed}/{len(face_files)} embeddings stored")


@app.post("/faces/index")
def build_face_index():
    if _face_index_status.get("running"):
        return {"status": "already_running", **_face_index_status}
    threading.Thread(target=_face_index_worker, daemon=True).start()
    return {"status": "started"}


@app.get("/faces/index/status")
def face_index_status():
    return _face_index_status


@app.get("/faces/{group_id}")
def face_status(group_id: str):
    """Report whether a face crop exists for a group.

    Face extraction runs asynchronously after "set preferred", so the studio
    polls this (over HTTP, which works even when the page is opened via file://)
    instead of guessing with an <img> load that 404s during the race window.
    """
    face_fname = f"{group_id.replace('/', '_')}_face.png"
    face_path = os.path.join(_load_output_dir(), face_fname)
    return {"exists": os.path.exists(face_path), "filename": face_fname}


@app.post("/images/{filename}/undress")
def undress_image(filename: str, background_tasks: BackgroundTasks):
    """Queue a generation using the undress prompt on the given image."""
    output_dir = _load_output_dir()
    fpath = os.path.join(output_dir, filename)
    if not os.path.exists(fpath):
        raise HTTPException(404, "Image not found")
    person = database.get_person(filename)
    group_id = person[1] if person and person[1] else naming.get_base_name(filename)
    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {"status": "queued", "done": 0, "failed": 0, "total": 1}
    threading.Thread(
        target=_batch_worker,
        args=(job_id, [filename], [UNDRESS_PROMPT], [None],
              random.randint(0, MAX_SEED), MAX_AREA),
        kwargs={"group_id": group_id},
        daemon=True,
    ).start()
    return {"job_id": job_id, "filename": filename}


@app.delete("/images/{filename}")
def delete_image(filename: str):
    person = database.get_person(filename)
    if person and person[5] and os.path.exists(person[5]):
        _move_to_trash(person[5])

    database.delete_person(filename)
    _invalidate_static()
    return {"status": "deleted", "filename": filename}


@app.delete("/groups/{group_id}")
def delete_group(group_id: str):
    files = database.get_group_files(group_id)
    for filename, filepath in files:
        if filepath and os.path.exists(filepath):
            _move_to_trash(filepath)

    database.delete_group(group_id)
    _invalidate_static()
    return {"status": "deleted", "group_id": group_id}


@app.post("/remove-background/{filename}")
def remove_background(filename: str):
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")

    path = person[5]
    with open(path, "rb") as f:
        png_bytes = f.read()

    transparent_png = _apply_transparency(png_bytes)

    with open(path, "wb") as f:
        f.write(transparent_png)

    # Persist the state + refresh static data so the flag (and No-BG/Crop buttons)
    # survive a page reload instead of reverting to has_background=True.
    database.upsert_person(filename, has_background=False)
    _invalidate_static()
    return {"status": "success", "filename": filename, "has_background": False}


@app.post("/images/{filename}/invert-alpha")
def invert_alpha(filename: str):
    """Invert the alpha channel in place — recovers cases where background removal
    kept the background and dropped the subject (the wrong segment)."""
    import numpy as np
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    img = Image.open(path).convert("RGBA")
    arr = np.array(img)
    arr[:, :, 3] = 255 - arr[:, :, 3]
    Image.fromarray(arr, "RGBA").save(path, format="PNG")
    database.upsert_person(filename, has_background=False)
    _invalidate_static()
    return {"status": "success", "filename": filename}


@app.post("/remove-background/group/{group_id}")
def remove_background_group(group_id: str, background_tasks: BackgroundTasks):
    def _bg_task():
        files = database.get_group_files(group_id)
        for filename, filepath in files:
            if filepath and os.path.exists(filepath):
                try:
                    with open(filepath, "rb") as f:
                        png_bytes = f.read()
                    transparent_png = _apply_transparency(png_bytes)
                    with open(filepath, "wb") as f:
                        f.write(transparent_png)
                except Exception as e:
                    print(f"Error removing background for {filename}: {e}")
    
    background_tasks.add_task(_bg_task)
    return {"status": "processing", "group_id": group_id}


# --- scenery generation -------------------------------------------------------

def _extract_frame_at(video_path: str, t: float) -> Image.Image:
    """Extract a single frame at time t (seconds) from a video via ffmpeg."""
    import subprocess as _sp
    r = _sp.run(
        ['ffmpeg', '-y', '-ss', str(t), '-i', video_path,
         '-frames:v', '1', '-f', 'image2pipe', '-vcodec', 'png', 'pipe:1'],
        capture_output=True, timeout=15,
    )
    if r.returncode != 0 or not r.stdout:
        raise ValueError(f"ffmpeg frame extract failed: {r.stderr.decode(errors='replace')[:300]}")
    return Image.open(io.BytesIO(r.stdout)).convert("RGB")


class SceneryRequest(BaseModel):
    model_filename: str           # person image in output_dir → image2 (Picture 2)
    scene_bytes: str | None = None  # base64-encoded PNG/JPEG of the reference scene → image1
    scene_video: str | None = None  # wireframe video name to extract frame from
    scene_time: float = 0.0       # timestamp (seconds) to extract from video
    extra_filename: str | None = None  # optional extra reference in output_dir → image3 (Picture 3)
    prompt: str | None = None     # override; auto-built if None
    seed: int = -1


def _make_side_by_side(img1: Image.Image, img2: Image.Image,
                       max_h: int = 1024) -> Image.Image:
    """Combine two images side by side at the same height (capped at max_h)."""
    target_h = min(max(img1.height, img2.height), max_h)
    r1 = target_h / img1.height
    r2 = target_h / img2.height
    w1, w2 = int(img1.width * r1), int(img2.width * r2)
    img1_r = img1.resize((w1, target_h), Image.LANCZOS)
    img2_r = img2.resize((w2, target_h), Image.LANCZOS)
    out = Image.new("RGB", (w1 + w2, target_h))
    out.paste(img1_r, (0, 0))
    out.paste(img2_r, (w1, 0))
    return out


def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
                    prompt: str, seed: int, extra_pils: list | None = None,
                    scene_video: str | None = None, extra_filename: str | None = None):
    output_dir = _load_output_dir()
    try:
        model_path = os.path.join(output_dir, model_filename)
        model_pil  = Image.open(model_path).convert("RGB")

        # image1=scene (→ Picture 1, output sized to scene), image2=person (→ Picture 2),
        # image3=optional extra ref (→ Picture 3). The node prepends "Picture 1: <img>
        # Picture 2: <img> ..." to the prompt so the model can reason about each by name.
        extra_images = [model_pil] + list(extra_pils or [])
        png_bytes = _run_pipeline(
            scene_pil.convert("RGB"), prompt, seed, MAX_AREA,
            extra_images=extra_images,
        )
        ts         = time.strftime("%Y%m%d_%H%M%S")
        base_name  = naming.get_base_name(model_filename)
        # Ensure .png extension
        if not base_name.lower().endswith('.png'):
            base_name = os.path.splitext(base_name)[0] + '.png'
        out_name   = f"{ts}_sc_{base_name}"
        out_path   = os.path.join(output_dir, out_name)
        with open(out_path, "wb") as f:
            f.write(png_bytes)
        person = database.get_person(model_filename)
        group_id = person[1] if person and person[1] else naming.get_base_name(model_filename)
        try:
            embedding  = embeddings.generate_embedding(out_path)
            next_order = database.get_next_sort_order(group_id)
            # Store all source references: person image, background video (if any), extra ref (if any)
            refs = [model_filename]
            if scene_video:
                refs.append(f"video:{scene_video}")
            if extra_filename:
                refs.append(extra_filename)
            database.upsert_person(out_name, filepath=out_path, embedding=embedding,
                                   group_id=group_id, prompt=prompt,
                                   sort_order=next_order,
                                   source_refs=json.dumps(refs))
        except Exception as db_err:
            print(f"[scenery] DB error: {db_err}")
        jobs[job_id]["status"] = "done"
        jobs[job_id]["output"] = out_name
        # Regenerate the static JSON so the frontend's polling surfaces the new
        # image immediately (matching _batch_worker / _multi_ref_worker). Without
        # this the output only appeared after a server restart.
        _invalidate_static()
    except Exception as e:
        print(f"[scenery] error: {e}")
        jobs[job_id]["status"] = "error"
        jobs[job_id]["error"]  = str(e)


@app.post("/generate-scenery")
def generate_scenery(req: SceneryRequest):
    output_dir    = _load_output_dir()
    wireframe_dir = _load_wireframe_dir()

    model_path = os.path.join(output_dir, req.model_filename)
    if not os.path.exists(model_path):
        raise HTTPException(404, f"Model image not found: {req.model_filename}")

    # Resolve scene image
    if req.scene_bytes:
        import base64
        raw = base64.b64decode(req.scene_bytes)
        scene_pil = Image.open(io.BytesIO(raw)).convert("RGB")
    elif req.scene_video:
        video_path = os.path.join(wireframe_dir, req.scene_video)
        if not os.path.exists(video_path):
            raise HTTPException(404, f"Scene video not found: {req.scene_video}")
        try:
            scene_pil = _extract_frame_at(video_path, req.scene_time)
        except Exception as e:
            raise HTTPException(500, f"Frame extraction failed: {e}")
    else:
        raise HTTPException(400, "Provide scene_bytes or scene_video")

    # Optional image3 (Picture 3) — extra reference from output_dir
    extra_pils = []
    if req.extra_filename:
        extra_path = os.path.join(output_dir, req.extra_filename)
        if not os.path.exists(extra_path):
            raise HTTPException(404, f"Extra reference not found: {req.extra_filename}")
        extra_pils.append(Image.open(extra_path).convert("RGB"))

    prompt = req.prompt or (
        "Place the person from Picture 2 naturally inside the environment shown in Picture 1. "
        "Keep the person's face, body proportions, clothing and pose exactly as in Picture 2. "
        "Use the location, lighting and atmosphere from Picture 1 as the background. "
        "Match the color temperature and shadows so it looks like one photograph taken on location. "
        + ("Incorporate the reference from Picture 3 as well. " if extra_pils else "")
        + "Output a single photorealistic image. High quality, detailed."
    )

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {"status": "running", "type": "scenery", "total": 1, "done": 0, "failed": 0}
    threading.Thread(
        target=_scenery_worker,
        args=(job_id, req.model_filename, scene_pil, prompt, req.seed, extra_pils),
        kwargs={"scene_video": req.scene_video if req.scene_video else None,
                "extra_filename": req.extra_filename},
        daemon=True,
    ).start()
    return {"job_id": job_id, "model": req.model_filename}


@app.get("/scenery/library")
def scenery_library():
    """Return all scenery images grouped by source video reference."""
    conn = database.get_db_connection()
    cur = conn.cursor()
    try:
        cur.execute("""
            SELECT filename, group_id, source_refs
            FROM person
            WHERE archived IS NOT TRUE
              AND filename LIKE '%_sc_%'
              AND source_refs IS NOT NULL
            ORDER BY filename DESC
        """)
        rows = cur.fetchall()
    finally:
        cur.close()
        database._put_db_connection(conn)

    by_video: dict[str, list] = {}
    ungrouped: list = []
    for filename, group_id, source_refs_raw in rows:
        try:
            refs = json.loads(source_refs_raw) if source_refs_raw else []
        except Exception:
            refs = []
        video = next((r[len("video:"):] for r in refs if r.startswith("video:")), None)
        entry = {"filename": filename, "group_id": group_id, "refs": refs}
        if video:
            by_video.setdefault(video, []).append(entry)
        else:
            ungrouped.append(entry)

    groups = [{"video": v, "items": items} for v, items in by_video.items()]
    return {"groups": groups, "ungrouped": ungrouped}


# --- SAM2 background removal --------------------------------------------------

_sam2_predictor      = None
_sam2_predictor_lock = threading.Lock()


def _load_sam2():
    """Lazy-load SAM2 image predictor. Returns predictor or False if unavailable."""
    global _sam2_predictor
    if _sam2_predictor is not None:
        return _sam2_predictor
    with _sam2_predictor_lock:
        if _sam2_predictor is not None:
            return _sam2_predictor
        try:
            from sam2.build_sam import build_sam2
            from sam2.sam2_image_predictor import SAM2ImagePredictor
            with open(CONFIG_PATH) as f:
                conf = json.load(f)
            ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
            cfg  = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
            if not os.path.exists(ckpt):
                raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")
            model = build_sam2(cfg, ckpt, device="cuda")
            _sam2_predictor = SAM2ImagePredictor(model)
            print(f"[sam2] loaded from {ckpt}")
        except Exception as e:
            print(f"[sam2] not available: {e}")
            _sam2_predictor = False
        return _sam2_predictor


def _person_mask_score(mask, h: int, w: int):
    """Rate how much `mask` looks like a centered subject vs. the background.

    The subject (person) sits in the middle of the frame and rarely fills the
    corners; the background is the opposite — it hugs the corners and is sparse
    in the center.  So `center_cov - corner_cov` is strongly positive for a
    correct person mask and negative when SAM2 has selected the background
    instead (the inverted-mask failure mode).

    Returns (score, center_cov, corner_cov), all floats in [-1, 1] / [0, 1].
    """
    import numpy as np
    cy0, cy1 = int(h * 0.30), int(h * 0.70)
    cx0, cx1 = int(w * 0.30), int(w * 0.70)
    center_cov = float(mask[cy0:cy1, cx0:cx1].mean())
    cs = max(1, int(min(h, w) * 0.10))
    corner_cov = float(np.mean([
        mask[:cs, :cs].mean(),  mask[:cs, -cs:].mean(),
        mask[-cs:, :cs].mean(), mask[-cs:, -cs:].mean(),
    ]))
    return center_cov - corner_cov, center_cov, corner_cov


def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
    """Remove background with SAM2 bbox segmentation; fallback to rembg.

    Uses a near-full-frame bbox so SAM2 finds the largest foreground object
    (the person) regardless of rotation or pose.  This works well because
    "transparent background" is stripped from the Qwen prompt upstream, so the
    model renders a solid real background — giving SAM2 clear contrast to work
    with.  Point prompts were tried but produced holes in ¾-rotated poses
    because the spine-column seeds land on background when the body is offset.
    """
    predictor = _load_sam2()
    if predictor is False:
        return _apply_transparency(png_bytes)
    try:
        import numpy as np
        import torch
        img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
        arr = np.array(img)
        h, w = arr.shape[:2]

        # Near-full-frame bbox — 1 % margin so hair / shoes are inside the hint.
        # SAM2 treats this as "find the prominent object within this region".
        box = np.array([[int(w * 0.01), int(h * 0.01),
                         int(w * 0.99), int(h * 0.99)]], dtype=np.float32)

        with torch.inference_mode():
            predictor.set_image(arr)
            masks, scores, _ = predictor.predict(
                box=box,
                multimask_output=True,
            )

        if masks is None or len(masks) == 0:
            print("[sam2] no masks returned, falling back to rembg")
            return _apply_transparency(png_bytes)

        # Pick the candidate that best matches a centered subject rather than
        # blindly trusting argmax(scores): on busy or low-contrast backgrounds
        # the top-confidence SAM2 mask is sometimes the background itself.
        # Combine the centered-subject prior with SAM2's own confidence.
        best = None
        best_rank = -1e9
        for i in range(len(masks)):
            m = masks[i].astype(bool)
            psc, _, _ = _person_mask_score(m, h, w)
            rank = psc + 0.10 * float(scores[i])
            if rank > best_rank:
                best_rank, best = rank, m

        # Inversion guard — the user's hint: the model is in the center.  If the
        # chosen mask still covers the corners more than the center, SAM2 picked
        # the background; flip the alpha so the person stays opaque.
        psc, ccov, kcov = _person_mask_score(best, h, w)
        if psc < 0:
            print(f"[sam2] mask inverted (center {ccov:.0%} < corners {kcov:.0%}) — flipping alpha")
            best = ~best

        # Sanity check: a person should cover 5 %–92 % of the frame
        coverage = float(best.sum()) / (h * w)
        if coverage < 0.05 or coverage > 0.92:
            print(f"[sam2] mask coverage {coverage:.1%} out of range, falling back to rembg")
            return _apply_transparency(png_bytes)

        mask_np = best.astype(np.uint8) * 255

        # Soft anti-aliased edge (radius 1 keeps accessory detail)
        try:
            from PIL import ImageFilter
            alpha_img = Image.fromarray(mask_np, mode="L")
            alpha_img = alpha_img.filter(ImageFilter.GaussianBlur(radius=1))
        except Exception:
            alpha_img = Image.fromarray(mask_np, mode="L")

        rgba = img.convert("RGBA")
        r, g, b, _ = rgba.split()
        out = Image.merge("RGBA", (r, g, b, alpha_img))
        buf = io.BytesIO()
        out.save(buf, format="PNG")
        print(f"[sam2] mask OK ({coverage:.1%} coverage)")
        return buf.getvalue()
    except Exception as e:
        print(f"[sam2] inference error, falling back to rembg: {e}")
        return _apply_transparency(png_bytes)


def _apply_transparency_black_bg(png_bytes: bytes) -> bytes:
    """Background removal for black-background Qwen output (bg_removal=sam2 mode).

    Strategy:
      1. Threshold: any pixel with max-channel > 25 is person (non-black).
         This correctly identifies the subject regardless of pose or rotation.
      2. Derive a tight person bounding-box from the threshold mask.
      3. Run SAM2 with that box for sub-pixel edge refinement.
         Accept SAM2 result only when its coverage is close (±30 pp) to the
         threshold estimate — this rejects the inverted-mask failure mode where
         SAM2 picks the large dark region as the "object".
      4. Fall back to the threshold mask (Gaussian-blurred edges) if SAM2
         is unavailable, errors, or diverges.

    Do NOT use the full-frame bbox here: on black-background images the large
    dark region scores higher than the person, causing SAM2 to invert the mask.
    """
    import numpy as np
    import torch
    from PIL import ImageFilter

    img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    arr = np.array(img)
    h, w = arr.shape[:2]

    # Step 1 — threshold: non-black pixels are the person
    is_person  = np.max(arr, axis=2) > 25
    thresh_cov = float(is_person.sum()) / (h * w)
    print(f"[bg-black] threshold coverage: {thresh_cov:.1%}")

    if not is_person.any():
        print("[bg-black] all-black image — falling back to rembg")
        return _apply_transparency(png_bytes)

    # Step 2 — tight bounding box from threshold
    rows = np.any(is_person, axis=1)
    cols = np.any(is_person, axis=0)
    rmin = int(np.where(rows)[0][0]);  rmax = int(np.where(rows)[0][-1])
    cmin = int(np.where(cols)[0][0]);  cmax = int(np.where(cols)[0][-1])
    margin = int(min(h, w) * 0.02)
    x1 = max(0, cmin - margin); y1 = max(0, rmin - margin)
    x2 = min(w, cmax + margin); y2 = min(h, rmax + margin)

    # Step 3 — SAM2 with the person-specific bbox
    predictor = _load_sam2()
    if predictor is not False:
        box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
        try:
            with torch.inference_mode():
                predictor.set_image(arr)
                masks, scores, _ = predictor.predict(box=box, multimask_output=True)

            if masks is not None and len(masks) > 0:
                best    = masks[int(np.argmax(scores))]
                sam_cov = float(best.sum()) / (h * w)
                print(f"[bg-black] SAM2 coverage: {sam_cov:.1%}")

                if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
                    mask_np   = best.astype(np.uint8) * 255
                    alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
                    rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
                    out  = Image.merge("RGBA", (r, g, b, alpha_img))
                    buf  = io.BytesIO(); out.save(buf, "PNG")
                    print(f"[bg-black] SAM2 accepted ✓")
                    return buf.getvalue()
                else:
                    print(f"[bg-black] SAM2 diverged ({sam_cov:.1%} vs {thresh_cov:.1%}) — threshold fallback")
        except Exception as e:
            print(f"[bg-black] SAM2 error: {e} — threshold fallback")

    # Step 4 — fallback: threshold mask with soft edge blur
    print("[bg-black] using threshold mask")
    mask_np   = is_person.astype(np.uint8) * 255
    alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
    rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
    out  = Image.merge("RGBA", (r, g, b, alpha_img))
    buf  = io.BytesIO(); out.save(buf, "PNG")
    return buf.getvalue()


@app.post("/remove-background-sam/{filename}")
def remove_background_sam(filename: str):
    """SAM2-based background removal.

    Writes the transparent result as a sidecar <stem>.nobg.png alongside the
    original, which is left untouched. Returns the sidecar URL so the UI can
    switch the viewer without touching the source file.
    Falls back to rembg when SAM2 is unavailable.
    """
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    output_dir = os.path.dirname(path)
    stem = os.path.splitext(filename)[0]
    nobg_filename = f"{stem}.nobg.png"
    nobg_path = os.path.join(output_dir, nobg_filename)

    with open(path, "rb") as f:
        png_bytes = f.read()
    transparent_png = _apply_transparency_sam2(png_bytes)
    with open(nobg_path, "wb") as f:
        f.write(transparent_png)

    # Register sidecar in DB so it appears in the same group
    group_id = person[1]
    database.upsert_person(nobg_filename, filepath=nobg_path,
                           group_id=group_id, has_background=False)

    used_sam2 = _sam2_predictor is not False and _sam2_predictor is not None
    return {
        "status": "success",
        "filename": filename,
        "nobg_filename": nobg_filename,
        "nobg_url": f"/output/{nobg_filename}",
        "used_sam2": used_sam2,
    }


@app.post("/images/{filename}/autocrop")
def autocrop_image(filename: str):
    """Crop away transparent borders from an image in-place."""
    import numpy as np
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    img = Image.open(path).convert("RGBA")
    arr = np.array(img)
    alpha = arr[:, :, 3]
    rows = np.any(alpha > 0, axis=1)
    cols = np.any(alpha > 0, axis=0)
    if not rows.any():
        raise HTTPException(400, "Image is fully transparent")
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    cropped = img.crop((cmin, rmin, cmax + 1, rmax + 1))
    cropped.save(path, format="PNG")
    return {"status": "success", "filename": filename, "box": [int(cmin), int(rmin), int(cmax+1), int(rmax+1)]}


class CropRequest(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    as_copy: bool = False  # True → crop a fresh copy, leaving the original untouched


@app.post("/images/{filename}/crop")
def manual_crop_image(filename: str, req: CropRequest):
    """Crop the image to the given pixel rectangle (in original image coordinates).

    By default the crop is applied in-place. When ``as_copy`` is set, a new copy is
    created first (referencing the original via ``source_refs``) and the crop is applied
    to that copy, so the original is preserved.
    """
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    src_path = person[5]

    if req.as_copy:
        # Mirror duplicate_image: copy file + register a DB row that points back to the original.
        from datetime import datetime as _dt
        output_dir = os.path.dirname(src_path)
        ext = os.path.splitext(filename)[1] or ".png"
        stem = os.path.splitext(filename)[0]
        ts = _dt.now().strftime("%Y%m%d_%H%M%S")
        new_filename = f"{ts}_crop_{stem}{ext}"
        path = os.path.join(output_dir, new_filename)
        shutil.copy2(src_path, path)
        database.upsert_person(
            new_filename, filepath=path, group_id=person[1],
            prompt=person[6], pose=person[7],
            has_background=person[11], has_clothing=person[13],
            source_refs=json.dumps([filename]),  # original is the reference
        )
    else:
        new_filename = filename
        path = src_path

    img = Image.open(path)
    w, h = img.size
    x1 = max(0, min(req.x1, w))
    y1 = max(0, min(req.y1, h))
    x2 = max(0, min(req.x2, w))
    y2 = max(0, min(req.y2, h))
    if x2 <= x1 or y2 <= y1:
        if req.as_copy:
            # Roll back the copy we just made so a bad rect doesn't leave an orphan.
            try:
                database.delete_person(new_filename)
                os.remove(path)
            except Exception:
                pass
        raise HTTPException(400, "Invalid crop rectangle")
    cropped = img.crop((x1, y1, x2, y2))
    fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
    cropped.save(path, format=fmt)
    if req.as_copy:
        _invalidate_static()
    return {"status": "success", "filename": filename, "new_filename": new_filename,
            "new_url": f"/output/{new_filename}", "as_copy": req.as_copy,
            "box": [x1, y1, x2, y2]}


class RotateRequest(BaseModel):
    degrees: int = 90  # clockwise rotation; must be a multiple of 90


@app.post("/images/{filename}/rotate")
def rotate_image(filename: str, req: RotateRequest):
    """Rotate an image clockwise in 90° steps, in place (lossless transpose)."""
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    deg = req.degrees % 360
    if deg not in (0, 90, 180, 270):
        raise HTTPException(400, "degrees must be a multiple of 90")
    if deg:
        # PIL transpose is defined counter-clockwise; map clockwise degrees onto it.
        cw_to_transpose = {
            90:  Image.Transpose.ROTATE_270,
            180: Image.Transpose.ROTATE_180,
            270: Image.Transpose.ROTATE_90,
        }
        img = Image.open(path).transpose(cw_to_transpose[deg])
        fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
        img.save(path, format=fmt)
    return {"status": "success", "filename": filename, "degrees": deg}


@app.post("/images/{filename}/duplicate")
def duplicate_image(filename: str):
    """Copy an image into the same group with a fresh timestamp-based filename."""
    import shutil as _shutil
    from datetime import datetime as _dt
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    output_dir = os.path.dirname(path)
    ext = os.path.splitext(filename)[1] or ".png"
    ts = _dt.now().strftime("%Y%m%d_%H%M%S")
    stem = os.path.splitext(filename)[0]
    new_filename = f"{ts}_dup_{stem}{ext}"
    new_path = os.path.join(output_dir, new_filename)
    _shutil.copy2(path, new_path)
    group_id = person[1]
    # person tuple: (name, group_id, tags, embedding, clip_description, filepath,
    #                prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing)
    database.upsert_person(
        new_filename, filepath=new_path, group_id=group_id,
        prompt=person[6], pose=person[7],
        has_background=person[11],
        has_clothing=person[13],
        source_refs=json.dumps([filename]),  # original is the reference
    )
    _invalidate_static()
    return {"status": "success", "new_filename": new_filename, "new_url": f"/output/{new_filename}"}


@app.post("/restore-background/{filename}")
def restore_background(filename: str):
    """Flatten RGBA → RGB (white composite), making the image opaque again."""
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    path = person[5]
    img = Image.open(path)
    if img.mode == "RGBA":
        bg = Image.new("RGB", img.size, (255, 255, 255))
        bg.paste(img, mask=img.split()[3])
        buf = io.BytesIO()
        bg.save(buf, format="PNG")
        with open(path, "wb") as f:
            f.write(buf.getvalue())
    return {"status": "success", "filename": filename}


@app.get("/sam2/check")
def sam2_check():
    """Return whether SAM2 is available."""
    predictor = _load_sam2()
    return {"sam2": predictor is not False and predictor is not None}


# --- 2D body-pose preview -----------------------------------------------------
# Estimates COCO-17 keypoints from the model image so the UI can overlay a
# posenet-style skeleton. Estimator is feature-detected: rtmlib (ONNX, reuses the
# already-installed onnxruntime) is preferred, mediapipe is a fallback. If neither
# is installed the endpoints report unavailable instead of erroring the request.

_pose_estimator = None          # cached (callable, backend_name) or False if unavailable
_pose_lock = threading.Lock()

# COCO-17 keypoint names (the order rtmlib's Body model returns).
POSE_KEYPOINT_NAMES = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]
# Bone connections (index pairs into COCO-17) for drawing the skeleton.
POSE_SKELETON = [
    (5, 7), (7, 9), (6, 8), (8, 10),          # arms
    (11, 13), (13, 15), (12, 14), (14, 16),   # legs
    (5, 6), (11, 12), (5, 11), (6, 12),       # torso
    (0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (0, 6),  # head/neck
]
# mediapipe Pose (33 landmarks) → COCO-17 index map.
_MP_TO_COCO = [0, 2, 5, 7, 8, 11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28]


def _load_pose_estimator():
    global _pose_estimator
    if _pose_estimator is not None:
        return _pose_estimator
    with _pose_lock:
        if _pose_estimator is not None:
            return _pose_estimator
        # Preferred: rtmlib (RTMPose, ONNX) — returns COCO-17 directly.
        try:
            from rtmlib import Body
            import numpy as np
            model = Body(mode="balanced", backend="onnxruntime", device="cpu")

            def _infer_rtm(pil):
                bgr = np.array(pil.convert("RGB"))[:, :, ::-1]
                kpts, scores = model(bgr)   # (N,17,2), (N,17)
                people = []
                for person_kpts, person_scores in zip(kpts, scores):
                    people.append([[float(x), float(y), float(s)]
                                   for (x, y), s in zip(person_kpts, person_scores)])
                return people

            _pose_estimator = (_infer_rtm, "rtmlib")
            print("[pose] using rtmlib (RTMPose)")
            return _pose_estimator
        except Exception as e:
            print(f"[pose] rtmlib unavailable: {e}")

        # Fallback: mediapipe Pose (single person, normalized landmarks).
        try:
            import mediapipe as mp
            import numpy as np
            mp_pose = mp.solutions.pose.Pose(static_image_mode=True, model_complexity=2)

            def _infer_mp(pil):
                rgb = np.array(pil.convert("RGB"))
                h, w = rgb.shape[:2]
                res = mp_pose.process(rgb)
                if not res.pose_landmarks:
                    return []
                lm = res.pose_landmarks.landmark
                kpts = []
                for mp_idx in _MP_TO_COCO:
                    p = lm[mp_idx]
                    kpts.append([float(p.x * w), float(p.y * h), float(p.visibility)])
                return [kpts]

            _pose_estimator = (_infer_mp, "mediapipe")
            print("[pose] using mediapipe Pose")
            return _pose_estimator
        except Exception as e:
            print(f"[pose] mediapipe unavailable: {e}")

        _pose_estimator = False
        return _pose_estimator


# --- pose similarity (descriptor + index) -------------------------------------
# Pose descriptors are normalized (translation + scale invariant) COCO-17 vectors,
# cached in <output>/_data/poses_index.json so we can rank library images by pose.

_POSE_MIN_SCORE = 0.3
# Left/right keypoint pairs for the mirror-invariant distance.
_POSE_MIRROR = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
_pose_index_status = {"running": False, "done": 0, "total": 0}


def _pose_descriptor(keypoints):
    """Normalize one person's COCO-17 keypoints into a translation/scale-invariant
    descriptor: {"vec": [34 floats], "vis": [17 ints]}. Returns None if too sparse."""
    vis = [1 if (kp[2] >= _POSE_MIN_SCORE) else 0 for kp in keypoints]
    if sum(vis) < 6:
        return None

    def _mid(a, b):
        if vis[a] and vis[b]:
            return ((keypoints[a][0] + keypoints[b][0]) / 2.0,
                    (keypoints[a][1] + keypoints[b][1]) / 2.0)
        return None

    hip = _mid(11, 12)
    sho = _mid(5, 6)
    center = hip or sho
    if center is None:
        return None
    # Scale by torso length; fall back to keypoint spread if torso isn't visible.
    if hip and sho:
        scale = ((hip[0] - sho[0]) ** 2 + (hip[1] - sho[1]) ** 2) ** 0.5
    else:
        xs = [keypoints[i][0] for i in range(17) if vis[i]]
        ys = [keypoints[i][1] for i in range(17) if vis[i]]
        scale = max(max(xs) - min(xs), max(ys) - min(ys))
    if not scale or scale < 1e-3:
        return None

    vec = []
    for i in range(17):
        if vis[i]:
            vec.append((keypoints[i][0] - center[0]) / scale)
            vec.append((keypoints[i][1] - center[1]) / scale)
        else:
            vec.extend([0.0, 0.0])
    return {"vec": vec, "vis": vis}


def _pose_distance(a, b):
    """Weighted L2 between two descriptors over jointly-visible joints, taking the
    min of the direct and left-right-mirrored comparison. Lower = more similar."""
    def _dist(av, avis, bv, bvis, mirror):
        total, n = 0.0, 0
        for i in range(17):
            j = _POSE_MIRROR[i] if mirror else i
            if not (avis[i] and bvis[j]):
                continue
            bx = bv[j * 2] * (-1 if mirror else 1)   # flip x when mirrored
            by = bv[j * 2 + 1]
            dx = av[i * 2] - bx
            dy = av[i * 2 + 1] - by
            total += dx * dx + dy * dy
            n += 1
        return (total / n) ** 0.5 if n >= 4 else float("inf")
    direct = _dist(a["vec"], a["vis"], b["vec"], b["vis"], False)
    mirror = _dist(a["vec"], a["vis"], b["vec"], b["vis"], True)
    return min(direct, mirror)


def _best_person(people):
    """Pick the largest-bbox person from an estimator result (most prominent subject)."""
    best, best_area = None, -1.0
    for kpts in people:
        xs = [k[0] for k in kpts if k[2] >= _POSE_MIN_SCORE]
        ys = [k[1] for k in kpts if k[2] >= _POSE_MIN_SCORE]
        if len(xs) < 2:
            continue
        area = (max(xs) - min(xs)) * (max(ys) - min(ys))
        if area > best_area:
            best, best_area = kpts, area
    return best


def _pose_index_path():
    return os.path.join(_load_output_dir(), "_data", "poses_index.json")


def _load_pose_index():
    try:
        with open(_pose_index_path(), "r") as f:
            return json.load(f)
    except Exception:
        return {}


_pose_index_lock = threading.Lock()


def _save_pose_index_entry(filename, desc):
    with _pose_index_lock:
        idx = _load_pose_index()
        idx[filename] = desc
        os.makedirs(os.path.dirname(_pose_index_path()), exist_ok=True)
        _write_json(_pose_index_path(), idx)


@app.get("/pose/check")
def pose_check():
    """Report whether a body-pose estimator is available (and which backend)."""
    est = _load_pose_estimator()
    if not est:
        return {"available": False,
                "hint": "pip install rtmlib onnxruntime  (or: pip install mediapipe)"}
    return {"available": True, "backend": est[1]}


@app.post("/images/{filename}/pose")
def estimate_pose(filename: str):
    """Estimate COCO-17 body keypoints for an image. Returns pixel-space keypoints
    plus the skeleton edge list so the frontend can overlay a posenet-style preview."""
    person = database.get_person(filename)
    if not person or not person[5] or not os.path.exists(person[5]):
        raise HTTPException(404, "Image file not found")
    est = _load_pose_estimator()
    if not est:
        raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
    infer, backend = est
    pil = Image.open(person[5]).convert("RGB")
    try:
        people = infer(pil)
    except Exception as e:
        raise HTTPException(500, f"Pose estimation failed: {e}")
    # Cache the descriptor so "find similar pose" can rank this image later.
    best = _best_person(people)
    if best is not None:
        desc = _pose_descriptor(best)
        if desc is not None:
            try:
                _save_pose_index_entry(filename, desc)
            except Exception as e:
                print(f"[pose] index save failed for {filename}: {e}")
    return {
        "status": "success",
        "backend": backend,
        "width": pil.width,
        "height": pil.height,
        "names": POSE_KEYPOINT_NAMES,
        "skeleton": POSE_SKELETON,
        "people": people,
    }


def _build_pose_index_task():
    try:
        est = _load_pose_estimator()
        if not est:
            return
        infer, _ = est
        output_dir = _load_output_dir()
        with _pose_index_lock:
            idx = _load_pose_index()
        persons = database.list_persons()
        todo = [p[0] for p in persons
                if p[0] not in idx
                and (p[12] or "image") != "video"
                and p[0].lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
        _pose_index_status.update(running=True, done=0, total=len(todo))
        print(f"[pose] index build: {len(todo)} images to process")
        dirty = 0
        for fn in todo:
            try:
                fpath = os.path.join(output_dir, fn)
                if os.path.exists(fpath):
                    best = _best_person(infer(Image.open(fpath).convert("RGB")))
                    desc = _pose_descriptor(best) if best is not None else None
                    if desc is not None:
                        idx[fn] = desc
                        dirty += 1
            except Exception as e:
                print(f"[pose] index error for {fn}: {e}")
            _pose_index_status["done"] += 1
            # Batch-flush every 50 to avoid O(n^2) full-file rewrites.
            if dirty >= 50:
                with _pose_index_lock:
                    _write_json(_pose_index_path(), idx)
                dirty = 0
                print(f"[pose] index progress: {_pose_index_status['done']}/{len(todo)}")
        with _pose_index_lock:
            _write_json(_pose_index_path(), idx)
        print(f"[pose] index build complete: {len(idx)} entries")
    except Exception as e:
        print(f"[pose] index build failed: {e}")
    finally:
        _pose_index_status["running"] = False


@app.post("/pose/index")
def build_pose_index():
    """Compute pose descriptors for all library images lacking one (daemon thread)."""
    if not _load_pose_estimator():
        raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
    if _pose_index_status.get("running"):
        return {"status": "already_running", **_pose_index_status}
    threading.Thread(target=_build_pose_index_task, daemon=True).start()
    return {"status": "started"}


@app.get("/pose/index/status")
def pose_index_status():
    idx = _load_pose_index()
    return {**_pose_index_status, "indexed": len(idx)}


def _rank_similar_poses(query_desc, limit, exclude=None):
    idx = _load_pose_index()
    scored = []
    for fn, desc in idx.items():
        if fn == exclude or not desc or "vec" not in desc:
            continue
        d = _pose_distance(query_desc, desc)
        if d != float("inf"):
            scored.append((d, fn))
    scored.sort(key=lambda x: x[0])
    groups = get_groups() if scored else {}
    return [{"filename": fn, "group_id": groups.get(fn), "distance": round(d, 4)}
            for d, fn in scored[:limit]]


@app.get("/pose/similar/{filename}")
def similar_pose(filename: str, limit: int = 12):
    """Rank library images by pose similarity to the given image."""
    idx = _load_pose_index()
    query = idx.get(filename)
    if query is None:
        # Compute on demand (also caches it).
        person = database.get_person(filename)
        if not person or not person[5] or not os.path.exists(person[5]):
            raise HTTPException(404, "Image file not found")
        est = _load_pose_estimator()
        if not est:
            raise HTTPException(501, "No pose estimator installed.")
        best = _best_person(est[0](Image.open(person[5]).convert("RGB")))
        query = _pose_descriptor(best) if best is not None else None
        if query is None:
            raise HTTPException(404, "No detectable pose in this image")
        try:
            _save_pose_index_entry(filename, query)
        except Exception:
            pass
    return {"filename": filename, "similar": _rank_similar_poses(query, limit, exclude=filename)}


class PoseSimilarRequest(BaseModel):
    keypoints: list[list[float]]   # [[x,y,score], ...17] in image pixels
    width: int = 0
    height: int = 0
    limit: int = 12


@app.post("/pose/similar")
def similar_pose_from_keypoints(req: PoseSimilarRequest):
    """Rank library images by similarity to a supplied (e.g. hand-edited) skeleton."""
    query = _pose_descriptor(req.keypoints)
    if query is None:
        raise HTTPException(400, "Supplied pose is too sparse to match")
    return {"similar": _rank_similar_poses(query, req.limit)}


# --- pose-guided generation ---------------------------------------------------
# Renders COCO-17 keypoints as an OpenPose-style skeleton image and passes it as
# image2 to Qwen so the VLM can match the body pose while preserving appearance.

# Color per POSE_SKELETON entry (18 bones in POSE_SKELETON order).
_OPENPOSE_BONE_COLORS = [
    (0,   80, 255),  # 5→7  L upper arm
    (0,  130, 255),  # 7→9  L lower arm
    (255,  50,   0), # 6→8  R upper arm
    (255, 110,   0), # 8→10 R lower arm
    (255, 200,   0), # 11→13 L upper leg
    (240, 240,   0), # 13→15 L lower leg
    (100, 255,   0), # 12→14 R upper leg
    ( 50, 220,   0), # 14→16 R lower leg
    (200,   0, 220), # 5→6  shoulders
    (200,   0, 220), # 11→12 hips
    (180,  60, 180), # 5→11 L torso side
    ( 60, 180, 180), # 6→12 R torso side
    (  0, 230, 230), # 0→1  nose–L eye
    (  0, 230, 230), # 0→2  nose–R eye
    (  0, 180, 180), # 1→3  L eye–ear
    (  0, 180, 180), # 2→4  R eye–ear
    ( 80,  80, 255), # 0→5  head–L shoulder
    (255,  80,  80), # 0→6  head–R shoulder
]

# Per-joint hue (HSV-like, converted to RGB for PIL):
def _joint_color(idx: int) -> tuple[int, int, int]:
    h = idx / 17.0 * 6.0
    x = 1 - abs(h % 2 - 1)
    if   h < 1: r, g, b = 1, x, 0
    elif h < 2: r, g, b = x, 1, 0
    elif h < 3: r, g, b = 0, 1, x
    elif h < 4: r, g, b = 0, x, 1
    elif h < 5: r, g, b = x, 0, 1
    else:       r, g, b = 1, 0, x
    return (int(r * 255), int(g * 255), int(b * 255))


def _render_openpose_image(keypoints: list, width: int, height: int) -> Image.Image:
    """Draw COCO-17 keypoints as a white stick-figure on a vivid magenta background.

    Vivid background (not black/dark) prevents Qwen from treating the skeleton as
    a real scene element and from blending the bone colors into the generated output.
    White lines on magenta read unambiguously as a technical reference diagram.
    """
    from PIL import ImageDraw
    MIN_CONF = 0.25
    cap_w = min(width,  896)
    cap_h = min(height, 1152)
    scale = min(cap_w / max(width, 1), cap_h / max(height, 1))
    out_w, out_h = int(width * scale), int(height * scale)

    img  = Image.new("RGB", (out_w, out_h), color=(200, 0, 180))   # vivid magenta
    draw = ImageDraw.Draw(img)

    lw = max(4, out_w // 80)   # line width scales with image
    jr = max(5, out_w // 70)   # joint radius

    kpts = keypoints  # [[x, y, conf], ...]

    # bones — white so they stand out on the magenta background without
    # adding any colors that Qwen might reproduce in the generated image
    for a, b in POSE_SKELETON:
        if a >= len(kpts) or b >= len(kpts):
            continue
        pa, pb = kpts[a], kpts[b]
        if pa[2] < MIN_CONF or pb[2] < MIN_CONF:
            continue
        ax, ay = int(pa[0] * scale), int(pa[1] * scale)
        bx, by = int(pb[0] * scale), int(pb[1] * scale)
        draw.line([(ax, ay), (bx, by)], fill=(255, 255, 255), width=lw)

    # joints — light yellow circles with dark outline for visibility
    for kp in kpts:
        if kp[2] < MIN_CONF:
            continue
        x, y = int(kp[0] * scale), int(kp[1] * scale)
        draw.ellipse([(x - jr, y - jr), (x + jr, y + jr)],
                     fill=(255, 240, 100), outline=(60, 0, 60), width=max(1, lw // 2))

    return img


class PoseFromWireframeRequest(BaseModel):
    video: str          # filename of a wireframe video in the output dir
    time_s: float = 0.0 # timestamp in seconds to extract the frame from


@app.post("/pose/from-wireframe")
def pose_from_wireframe(req: PoseFromWireframeRequest):
    """Extract a frame from a wireframe video, run pose estimation, return COCO-17 keypoints.

    This is the pose-discovery tool: scrub through a wireframe video to find a frame
    that shows the desired pose, then load those keypoints into the pose overlay without
    ever passing the wireframe image to Qwen.
    """
    import tempfile
    import subprocess as sp

    output_dir = _load_output_dir()
    video_path = os.path.join(output_dir, req.video)
    if not os.path.exists(video_path):
        raise HTTPException(404, f"Video not found: {req.video}")

    est = _load_pose_estimator()
    if not est:
        raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
    infer, backend = est

    # Extract one frame at the requested timestamp with ffmpeg.
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tf:
        frame_path = tf.name
    try:
        result = sp.run(
            [
                "/usr/bin/ffmpeg", "-y",
                "-ss", str(req.time_s),
                "-i", video_path,
                "-frames:v", "1",
                "-q:v", "2",
                frame_path,
            ],
            capture_output=True, timeout=30,
        )
        if result.returncode != 0:
            raise HTTPException(500, f"ffmpeg frame extract failed: {result.stderr.decode()[:300]}")

        pil = Image.open(frame_path).convert("RGB")
    finally:
        try:
            os.unlink(frame_path)
        except OSError:
            pass

    try:
        people = infer(pil)
    except Exception as e:
        raise HTTPException(500, f"Pose estimation failed: {e}")

    best = _best_person(people)
    if best is None:
        raise HTTPException(422, "No person detected in the extracted frame")

    return {
        "status": "success",
        "backend": backend,
        "width": pil.width,
        "height": pil.height,
        "names": POSE_KEYPOINT_NAMES,
        "skeleton": POSE_SKELETON,
        "keypoints": best,
        "time_s": req.time_s,
    }


class PoseGenRequest(BaseModel):
    model_filename: str
    keypoints: list[list[float]]   # 17 × [x, y, conf] in original image pixels
    width: int                     # original image dimensions the keypoints came from
    height: int
    gesture_name: str | None = None   # preset name for prompt enrichment and tagging
    prompt: str | None = None         # full override; auto-built if None
    seed: int = -1
    extra_filename: str | None = None # optional image3 extra reference


def _keypoints_to_pose_text(kpts: list, width: int, height: int) -> str:
    """Convert COCO-17 keypoints to a natural-language pose description for Qwen."""
    MIN_CONF = 0.25
    W, H = max(width, 1), max(height, 1)

    def vis(idx):
        return idx < len(kpts) and kpts[idx][2] >= MIN_CONF

    def pt(idx):
        return (kpts[idx][0] / W, kpts[idx][1] / H) if vis(idx) else None

    parts = []

    # --- head orientation ---
    nose, l_ear, r_ear = pt(0), pt(3), pt(4)
    if l_ear and r_ear:
        ear_mid_x = (l_ear[0] + r_ear[0]) / 2
        if nose:
            dx = nose[0] - ear_mid_x
            if dx < -0.05:
                parts.append("head turned to the left")
            elif dx > 0.05:
                parts.append("head turned to the right")
            else:
                parts.append("head facing forward")
    elif l_ear and not r_ear:
        parts.append("head turned strongly to the right")
    elif r_ear and not l_ear:
        parts.append("head turned strongly to the left")

    # --- body rotation (shoulders vs hips) ---
    l_sh, r_sh = pt(5), pt(6)
    l_hip, r_hip = pt(11), pt(12)
    if l_sh and r_sh:
        sh_w = abs(r_sh[0] - l_sh[0])
        if l_hip and r_hip:
            hip_w = abs(r_hip[0] - l_hip[0])
            ratio = sh_w / max(hip_w, 0.01)
            if ratio < 0.5:
                parts.append("body rotated ~90° (side profile)")
            elif ratio < 0.75:
                parts.append("body rotated ~45° (three-quarter view)")
            else:
                parts.append("body facing forward")

    # --- arm positions ---
    for side, sh_i, el_i, wr_i in [("left", 5, 7, 9), ("right", 6, 8, 10)]:
        sh, el, wr = pt(sh_i), pt(el_i), pt(wr_i)
        if not sh:
            continue
        if wr:
            dy = sh[1] - wr[1]   # positive = wrist above shoulder
            dx = abs(wr[0] - sh[0])
            if dy > 0.15:
                if el and pt(el_i):
                    el_dy = sh[1] - el[1]
                    if el_dy > 0.05:
                        parts.append(f"{side} arm raised straight up")
                    else:
                        parts.append(f"{side} arm raised with elbow bent")
                else:
                    parts.append(f"{side} arm raised above shoulder")
            elif dx > 0.25:
                parts.append(f"{side} arm extended sideways")
            elif el:
                el_dx = abs(el[0] - sh[0])
                if el_dx > 0.15 and abs(wr[1] - sh[1]) < 0.1:
                    parts.append(f"{side} arm extended forward")
                else:
                    parts.append(f"{side} arm at side")
            else:
                parts.append(f"{side} arm at side")
        elif el:
            el_dy = sh[1] - el[1]
            if el_dy > 0.1:
                parts.append(f"{side} elbow raised")

    # --- leg positions ---
    for side, hip_i, kn_i, an_i in [("left", 11, 13, 15), ("right", 12, 14, 16)]:
        hip, kn, an = pt(hip_i), pt(kn_i), pt(an_i)
        if not hip:
            continue
        if an:
            dx = abs(an[0] - hip[0])
            dy = an[1] - hip[1]   # positive = ankle below hip (normal standing)
            if dy < 0.1:
                parts.append(f"{side} leg raised / knee up")
            elif dx > 0.2:
                parts.append(f"{side} leg stepped out to the side")
            elif kn:
                kn_dx = abs(kn[0] - hip[0])
                if kn_dx > 0.15:
                    parts.append(f"{side} leg bent / lunging")
                else:
                    parts.append(f"{side} leg straight, standing")
            else:
                parts.append(f"{side} leg straight")

    if not parts:
        return "standing neutral pose"
    return ", ".join(parts)


def _pad_for_pose(pil: Image.Image, keypoints: list, margin: float = 0.15) -> tuple[Image.Image, tuple[int,int,int,int]]:
    """Expand the image so target keypoints fit within the canvas.

    Only expands when keypoints are actually outside [0,W]×[0,H]. Expansion is
    symmetric on each axis (equal left/right, equal top/bottom) so the person
    stays centred on the new canvas. Fill is black to match PIL's RGBA→RGB fill.
    """
    W, H = pil.size
    all_kps = [kp for kp in keypoints if len(kp) >= 2]
    if not all_kps:
        return pil, (0, 0, 0, 0)

    xs = [kp[0] for kp in all_kps]
    ys = [kp[1] for kp in all_kps]

    # How far each side actually exceeds the image bounds (0 if within bounds).
    over_l = max(0.0, -min(xs))
    over_r = max(0.0,  max(xs) - W)
    over_t = max(0.0, -min(ys))
    over_b = max(0.0,  max(ys) - H)

    # If nothing is outside, no padding needed at all.
    if over_l == over_r == over_t == over_b == 0:
        return pil, (0, 0, 0, 0)

    # Symmetric expansion per axis: take the larger overhang and add margin.
    px = int(max(over_l, over_r) + margin * W)
    py = int(max(over_t, over_b) + margin * H)

    new_w, new_h = W + 2 * px, H + 2 * py
    canvas = Image.new("RGB", (new_w, new_h), color=(0, 0, 0))
    canvas.paste(pil, (px, py))
    print(f"[pose-gen] padded {W}×{H} → {new_w}×{new_h}  (±{px}px x, ±{py}px y)")
    return canvas, (px, py, px, py)


def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
                     keypoints: list | None = None,
                     gesture_name: str | None = None,
                     extra_filename: str | None = None):
    output_dir = _load_output_dir()
    try:
        model_path = os.path.join(output_dir, model_filename)
        model_pil  = Image.open(model_path).convert("RGB")

        # Auto-pad if the target pose would extend outside the current image bounds.
        if keypoints:
            model_pil, _padding = _pad_for_pose(model_pil, keypoints)

        # No skeleton image — pose is encoded entirely in the text prompt.
        extra_images = None
        if extra_filename:
            ep = os.path.join(output_dir, extra_filename)
            extra_images = [Image.open(ep).convert("RGB")]

        png_bytes  = _run_pipeline(model_pil, prompt, seed, MAX_AREA, extra_images=extra_images)
        ts         = time.strftime("%Y%m%d_%H%M%S")
        base_name  = naming.get_base_name(model_filename)
        if not base_name.lower().endswith(".png"):
            base_name = os.path.splitext(base_name)[0] + ".png"
        out_name   = f"{ts}_pose_{base_name}"
        out_path   = os.path.join(output_dir, out_name)
        with open(out_path, "wb") as f:
            f.write(png_bytes)

        person   = database.get_person(model_filename)
        group_id = person[1] if person and person[1] else naming.get_base_name(model_filename)

        embedding  = embeddings.generate_embedding(out_path)
        next_order = database.get_next_sort_order(group_id)
        refs       = [model_filename]
        if gesture_name:
            refs.append(f"gesture:{gesture_name}")
        database.upsert_person(
            out_name, filepath=out_path, embedding=embedding,
            group_id=group_id, prompt=prompt,
            pose=gesture_name or "custom",
            sort_order=next_order,
            source_refs=json.dumps(refs),
        )

        jobs[job_id]["status"] = "done"
        jobs[job_id]["output"] = out_name
        _invalidate_static()
    except Exception as e:
        print(f"[pose-gen] error: {e}")
        jobs[job_id]["status"] = "error"
        jobs[job_id]["error"]  = str(e)


@app.post("/pose/render")
def render_pose_image(req: PoseSimilarRequest):
    """Render the supplied keypoints as an OpenPose-style skeleton PNG (base64) — for preview only."""
    import base64
    skeleton = _render_openpose_image(req.keypoints, req.width or 512, req.height or 768)
    buf = io.BytesIO()
    skeleton.save(buf, format="PNG")
    return {"image": base64.b64encode(buf.getvalue()).decode()}


@app.post("/generate-with-pose")
def generate_with_pose(req: PoseGenRequest):
    """Generate the person in a specific body pose. Pose is encoded as text — no skeleton image
    is passed to Qwen to avoid wireframe bleed-through in the output."""
    output_dir = _load_output_dir()
    model_path = os.path.join(output_dir, req.model_filename)
    if not os.path.exists(model_path):
        raise HTTPException(404, f"Model image not found: {req.model_filename}")
    if req.extra_filename:
        ep = os.path.join(output_dir, req.extra_filename)
        if not os.path.exists(ep):
            raise HTTPException(404, f"Extra reference not found: {req.extra_filename}")

    if req.prompt:
        prompt = req.prompt
    else:
        # Convert skeleton keypoints to natural-language description
        pose_text = _keypoints_to_pose_text(req.keypoints, req.width, req.height)
        if req.gesture_name:
            pose_text = f'{req.gesture_name} gesture: {pose_text}'
        prompt = (
            f"Change the body pose of the person in image 1 to: {pose_text}. "
            "Keep the person's face, hair, skin tone, clothing, and body proportions "
            "exactly as they are — only the limb positions and body orientation change. "
            "Transparent background."
        )

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {"status": "running", "type": "pose_gen", "total": 1, "done": 0, "failed": 0}
    threading.Thread(
        target=_pose_gen_worker,
        args=(job_id, req.model_filename, prompt, req.seed),
        kwargs={
            "keypoints": req.keypoints,
            "gesture_name": req.gesture_name,
            "extra_filename": req.extra_filename,
        },
        daemon=True,
    ).start()
    return {"job_id": job_id, "model": req.model_filename}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),
                port=int(os.environ.get("PORT", "8500")))