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

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

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


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


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


# --- 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 _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)
                        database.upsert_person(
                            out_name, filepath=out_path, embedding=embedding,
                            group_id=actual_gid, prompt=prompt, pose=pose,
                            has_background=has_bg,
                            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)
                database.upsert_person(out_name, filepath=out_path, embedding=embedding,
                                       group_id=output_gid, prompt=prompt, pose=pose,
                                       has_background=has_bg,
                                       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()
    extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg')
    try:
        # Try to get from DB first
        try:
            persons = database.list_persons()
            # persons: (filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs)
            db_images = []
            for p in persons:
                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],
                })
            db_images.sort(
                key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"]))
                              if os.path.exists(os.path.join(output_dir, x["filename"])) else 0,
                reverse=True,
            )
            return {"images": db_images}
        except Exception as db_err:
            print(f"DB error in list_images: {db_err}")
            # Fallback to filesystem
            files = [f for f in os.listdir(output_dir) if f.lower().endswith(extensions)]
            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))


# --- 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)
        auto_name = 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)
        except Exception as db_err:
            print(f"Database error during tag: {db_err}")

        return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30]}
    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.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)
        auto_name = name or naming.generate_associative_name(tags)
        
        # 2. Embedding for source
        embedding = embeddings.generate_embedding(file_path)
        
        # 3. Register source in DB (optional, but good for tracking)
        # We'll use the original filename or a timestamped one
        database.upsert_person(filename, filepath=file_path, name=auto_name, clip_description=clip_desc, tags=tags, embedding=embedding, group_id=group_id)
        
        # 4. Default behavior: Crop if needed
        # We'll use default values from watcher
        cropped_pil = _crop_to_bbox(pil)
        
        # 5. Run prompts
        for i, prompt in enumerate(prompts):
            try:
                # Use RGB for pipeline
                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)
                
                # Register output in DB
                out_embedding = embeddings.generate_embedding(out_path)
                database.upsert_person(out_name, filepath=out_path, name=auto_name, clip_description=clip_desc, embedding=out_embedding, group_id=group_id)
            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.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}


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