"""
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}")


# --- 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()

_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):
    """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
        width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        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)
        tmp_name  = f"{ts}_fs_tmp_{vid_stem}_{base_name}.mp4"
        out_name  = f"{ts}_fs_{vid_stem}_{base_name}.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:
                result = frame.copy()
                for face in tgt_faces:
                    try:
                        result = swapper.get(result, 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
            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]),
        )

        jobs[job_id]["status"] = "done"
        jobs[job_id]["output"] = out_name

    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):
    """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)
    out_name   = f'{ts}_fs_{vid_stem}_{base_name}.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',
        '--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',
        '--face-selector-mode', 'many',
    ]
    if enhance:
        cmd += ['--face-enhancer-model', 'gfpgan_1.4']

    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',
        )
        # Read stdout for progress, stderr for error info
        import threading as _thr
        def _drain_stderr():
            for ln in proc.stderr:
                output_lines.append(ln.rstrip())
                print(f'[facefusion] {ln.rstrip()}')
        _thr.Thread(target=_drain_stderr, daemon=True).start()
        for line in proc.stdout:
            print(f'[facefusion] {line.rstrip()}')
            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
        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]),
        )
        jobs[job_id]['status'] = 'done'
        jobs[job_id]['output'] = out_name

    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 _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]

    # Transparency detection
    is_transparent = any(kw in prompt.lower() for kw in ["transparent", "no background", "remove background", "alpha channel"])
    if is_transparent:
        graph[NODE_NEGATIVE]["inputs"]["prompt"] = "checkerboard, grid, pattern, texture, background details, watermark, deformed anatomy"

    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 is_transparent:
        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}")


def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
                  seed: int, max_area: int, group_id: str | None = None):
    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")
            for prompt, pose in zip(prompts, poses):
                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)

                    png = _run_pipeline(pil, prompt, seed, max_area)
                    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
                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"


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
        except Exception as e:
            print(f"Error in multi-ref for prompt '{prompt}': {e}")
            jobs[job_id]["failed"] += 1

    jobs[job_id]["status"] = "done"


# --- 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)
    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


@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}
    t = threading.Thread(
        target=_batch_worker,
        args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
        daemon=True,
    )
    t.start()
    return {"job_id": job_id, "total": total_tasks}


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()


@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():
    output_dir = _load_output_dir()
    all_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg') + VIDEO_EXTENSIONS
    try:
        try:
            persons = database.list_persons()
            # 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
            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],
                })
            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/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 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)


@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},
            daemon=True,
        )
    else:
        t = threading.Thread(
            target=_faceswap_worker,
            args=(job_id, req.model_filename, req.video_name),
            kwargs={'enhance': req.enhance},
            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}")

    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}")

    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}")

    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))
    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))
    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])
    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 _process_upload(file_path: str, filename: str, prompts: list[str], name: str | None = None, group_id: 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,
        )

        # 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)
                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,
                )
            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}")


@app.post("/upload")
def upload_image(
    background_tasks: BackgroundTasks,
    image: UploadFile = File(...),
    prompts: str = Form(""),
    name: str = Form(None),
):
    # 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)
    # 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)
    
    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")]
    
    group_id = f"up_{uuid.uuid4().hex[:8]}"  # unique per upload; avoids collisions when pasting generic filenames
    background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, group_id)
    
    return {"status": "processing", "filename": filename, "group_id": group_id, "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))
    return {"filename": filename, "hidden": hidden}


@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)
    return {"filename": filename, "group_id": group_id}


@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)
    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)
    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)
        
    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
    scene_bytes: str | None = None  # base64-encoded PNG/JPEG of the reference scene
    scene_video: str | None = None  # wireframe video name to extract frame from
    scene_time: float = 0.0       # timestamp (seconds) to extract from video
    prompt: str | None = None     # override; auto-built if None
    seed: int = -1


def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
                    prompt: str, seed: int):
    output_dir = _load_output_dir()
    try:
        model_path = os.path.join(output_dir, model_filename)
        model_pil  = Image.open(model_path).convert("RGB")
        png_bytes  = _run_pipeline(model_pil, prompt, seed, MAX_AREA,
                                   extra_images=[scene_pil])
        ts         = time.strftime("%Y%m%d_%H%M%S")
        base_name  = naming.get_base_name(model_filename)
        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)
            database.upsert_person(out_name, filepath=out_path, embedding=embedding,
                                   group_id=group_id, prompt=prompt,
                                   sort_order=next_order,
                                   source_refs=json.dumps([model_filename]))
        except Exception as db_err:
            print(f"[scenery] DB error: {db_err}")
        jobs[job_id]["status"] = "done"
        jobs[job_id]["output"] = out_name
    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")

    prompt = req.prompt or (
        "Place this person into the provided scene. Keep the person's appearance, "
        "lighting, and pose natural and consistent with the environment. "
        "Photorealistic, high quality."
    )

    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),
        daemon=True,
    ).start()
    return {"job_id": job_id, "model": req.model_filename}


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

_sam2_predictor      = None
_sam2_predictor_lock = threading.Lock()


def _load_sam2():
    """Lazy-load SAM2 AutomaticMaskGenerator. Returns generator 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.automatic_mask_generator import SAM2AutomaticMaskGenerator
            with open(CONFIG_PATH) as f:
                conf = json.load(f)
            ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam2/sam2.1_hiera_tiny.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 = SAM2AutomaticMaskGenerator(model)
            print(f"[sam2] loaded from {ckpt}")
        except Exception as e:
            print(f"[sam2] not available: {e}")
            _sam2_predictor = False
        return _sam2_predictor


def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
    """Remove background using SAM2 (largest-area mask = subject), fallback to rembg."""
    predictor = _load_sam2()
    if predictor is False:
        return _apply_transparency(png_bytes)
    try:
        import numpy as np
        img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
        arr = np.array(img)
        masks = predictor.generate(arr)
        if not masks:
            return _apply_transparency(png_bytes)
        best = max(masks, key=lambda m: m["area"])
        mask_np = (best["segmentation"].astype(np.uint8) * 255)
        rgba = img.convert("RGBA")
        r, g, b, _ = rgba.split()
        alpha = Image.fromarray(mask_np, mode="L")
        out = Image.merge("RGBA", (r, g, b, alpha))
        buf = io.BytesIO()
        out.save(buf, format="PNG")
        return buf.getvalue()
    except Exception as e:
        print(f"[sam2] inference error, falling back to rembg: {e}")
        return _apply_transparency(png_bytes)


@app.post("/remove-background-sam/{filename}")
def remove_background_sam(filename: str):
    """SAM2-based background removal (RGBA PNG). Falls back to rembg if SAM2 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]
    with open(path, "rb") as f:
        png_bytes = f.read()
    transparent_png = _apply_transparency_sam2(png_bytes)
    with open(path, "wb") as f:
        f.write(transparent_png)
    used_sam2 = _sam2_predictor is not False and _sam2_predictor is not None
    return {"status": "success", "filename": filename, "used_sam2": used_sam2}


@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}


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