"""
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")
NAMES_PATH   = os.path.join(os.path.dirname(os.path.abspath(__file__)), "names.json")
GROUPS_PATH  = os.path.join(os.path.dirname(os.path.abspath(__file__)), "groups.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

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=["*"],
)


# --- 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 _load_json(path: str) -> dict:
    if os.path.exists(path):
        with open(path) as f:
            return json.load(f)
    return {}


def _save_json(path: str, data: dict):
    with open(path, "w") as f:
        json.dump(data, f, indent=2)


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_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] = " ".join(current_desc).strip()
                current_pose = line[2:].rstrip(":").strip()
                current_desc = []
            elif line and current_pose:
                current_desc.append(line)
        
        if current_pose:
            poses[current_pose] = " ".join(current_desc).strip()
            
    return poses


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",
) -> 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
    
    # 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"
        
    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], seed: int, max_area: int, group_id: str | None = None):
    output_dir = _load_output_dir()
    for fname in filenames:
        # If no group_id provided, try to inherit from source
        actual_gid = group_id
        if not actual_gid:
            try:
                person = database.get_person(fname)
                if person and person[1]: # group_id is at index 1
                    actual_gid = person[1]
                else:
                    # Create a new group for this standalone image
                    actual_gid = naming.get_base_name(fname)
                    # Update source image to join this group
                    database.upsert_person(fname, group_id=actual_gid)
            except Exception as e:
                print(f"Error determining group for {fname}: {e}")

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

        try:
            pil = Image.open(fpath).convert("RGB")
            for prompt in prompts:
                try:
                    png = _run_pipeline(pil, prompt, seed, max_area)
                    ts = time.strftime("%Y%m%d_%H%M%S")
                    
                    # Clean filename to avoid nested timestamps
                    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)
                    
                    # Register in DB
                    try:
                        embedding = embeddings.generate_embedding(out_path)
                        database.upsert_person(out_name, filepath=out_path, embedding=embedding, group_id=actual_gid)
                    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"


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


@app.post("/batch")
def start_batch(req: BatchRequest):
    prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
    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, req.seed, req.max_area, req.group_id),
        daemon=True,
    )
    t.start()
    return {"job_id": job_id, "total": total_tasks}


@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 is (filename, name, group_id, clip_description)
            db_images = []
            for p in persons:
                db_images.append({
                    "filename": p[0],
                    "name": p[1],
                    "group_id": p[2],
                    "clip_description": p[3]
                })
            # Still sort by mtime for consistency with filesystem
            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}")
            
        # Legacy fallback
        try:
            names = _load_json(NAMES_PATH)
            names[req.filename] = clip_desc
            _save_json(NAMES_PATH, names)
        except: pass
        
        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:
        return _load_json(NAMES_PATH)


@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}")
    
    # Legacy fallback
    try:
        names = _load_json(NAMES_PATH)
        names[filename] = name
        _save_json(NAMES_PATH, names)
    except: pass
    
    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:
        return _load_json(GROUPS_PATH)


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}")
    
    # Legacy fallback
    try:
        groups = _load_json(GROUPS_PATH)
        for fname in req.filenames:
            groups[fname] = gid
        _save_json(GROUPS_PATH, groups)
    except: pass
    
    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}")
    
    # Legacy fallback
    try:
        groups = _load_json(GROUPS_PATH)
        groups[req.filename] = gid
        _save_json(GROUPS_PATH, groups)
    except: pass
    
    return {"filename": req.filename}


@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 = naming.get_base_name(filename)
    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.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")))