"""
HDRemover.com — Complete SaaS Background Removal
HuggingFace Space | Gradio 4.x
"""

import os, io, time, json, logging, random, string, threading
from datetime import datetime, timezone, timedelta
from collections import deque
import numpy as np
from PIL import Image

logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
logger = logging.getLogger("hdremover")

def utcnow(): return datetime.now(timezone.utc)

# ── Config ─────────────────────────────────────────────────────────────────────
HF_TOKEN   = os.environ.get("HF_TOKEN", "")
MASTER_KEY = os.environ.get("MASTER_API_KEY", "hdremover-master-key-change-me")
ADMIN_PASS = os.environ.get("ADMIN_PASSWORD", "hdremover-admin-2026")
HF_REPO_ID = os.environ.get("HF_REPO_ID", "hdremover/background-remover")
HF_DB_REPO = os.environ.get("HF_DB_REPO", "hdremover/background-remover-db")
DB_FILE    = "api_keys.json"

# ── HF Login ───────────────────────────────────────────────────────────────────
if HF_TOKEN:
    try:
        from huggingface_hub import login
        login(token=HF_TOKEN, add_to_git_credential=False)
        logger.info("✅ HF login OK")
    except Exception as e:
        logger.warning(f"HF login: {e}")

# ── Dataset Persistence ────────────────────────────────────────────────────────
_flush_counter = 0
_last_flush_ts = 0.0
FLUSH_EVERY_N  = 50
FLUSH_EVERY_S  = 300

def load_keys_from_dataset() -> dict:
    if not HF_TOKEN: return {}
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=HF_DB_REPO, filename=DB_FILE,
                               repo_type="dataset", token=HF_TOKEN, force_download=True)
        with open(path) as f: data = json.load(f)
        logger.info(f"✅ Dataset loaded — {len(data)} keys")
        return data
    except Exception as e:
        logger.warning(f"Dataset load failed: {e}")
        return {}

def save_keys_to_dataset(force=False) -> bool:
    global _flush_counter, _last_flush_ts
    if not HF_TOKEN: return False
    now = utcnow().timestamp()
    if not force:
        if _flush_counter < FLUSH_EVERY_N and (now - _last_flush_ts) < FLUSH_EVERY_S:
            return False
    try:
        from huggingface_hub import upload_file
        safe = {}
        for k, v in API_KEYS.items():
            entry = dict(v)
            if "minute_window" in entry:
                entry["minute_window"] = list(entry["minute_window"])
            safe[k] = entry
        content = json.dumps(safe, indent=2, default=str).encode("utf-8")
        upload_file(path_or_fileobj=io.BytesIO(content), path_in_repo=DB_FILE,
                    repo_id=HF_DB_REPO, repo_type="dataset", token=HF_TOKEN,
                    commit_message=f"flush {len(API_KEYS)} keys")
        _flush_counter = 0; _last_flush_ts = now
        logger.info(f"✅ Dataset saved — {len(API_KEYS)} keys")
        return True
    except Exception as e:
        logger.error(f"Dataset save failed: {e}")
        return False

def _async_save(force=False):
    threading.Thread(target=save_keys_to_dataset, args=(force,), daemon=True).start()

# ── API Keys ───────────────────────────────────────────────────────────────────
API_KEYS: dict = {}
_from_dataset = load_keys_from_dataset()
if _from_dataset:
    API_KEYS = _from_dataset
    for v in API_KEYS.values():
        if "minute_window" in v:
            v["minute_window"] = deque(v["minute_window"])
    logger.info(f"✅ Loaded {len(API_KEYS)} keys from dataset")
else:
    _raw = os.environ.get("API_KEYS_JSON", "")
    if _raw:
        try:
            API_KEYS = json.loads(_raw)
            logger.info(f"✅ Loaded {len(API_KEYS)} keys from HF Secret")
            _async_save(force=True)
        except Exception as e:
            logger.warning(f"API_KEYS_JSON error: {e}")

API_KEYS.setdefault(MASTER_KEY, {
    "plan": "master", "calls_today": 0,
    "reset_at": (utcnow() + timedelta(days=1)).timestamp(),
    "owner": "HDRemover.com", "created_at": utcnow().isoformat()
})

PLAN_LIMITS = {
    "free":    {"daily": 10,     "per_minute": 2,      "models": ["fast"]},
    "starter": {"daily": 100,    "per_minute": 10,     "models": ["fast","quality","hair"]},
    "pro":     {"daily": 500,    "per_minute": 30,     "models": ["fast","quality","best","hair"]},
    "master":  {"daily": 999999, "per_minute": 999999, "models": ["fast","quality","best","hair"]},
}
PLAN_DEFAULT_EXPIRY = {"free": 30, "starter": 365, "pro": 365, "master": None}

def push_keys_to_hf_secret() -> tuple:
    _async_save(force=True)
    return True, "✅ **Saved!** Active immediately."

def gen_key(plan: str) -> str:
    chars = string.ascii_lowercase + string.digits
    r = lambda n: ''.join(random.choices(chars, k=n))
    return f"hdremover-{plan[:2]}-{r(8)}-{r(8)}"

# ── Constants ──────────────────────────────────────────────────────────────────
MAX_IMAGE_BYTES = 10 * 1024 * 1024
MAX_IMAGE_PX    = 4000

# ── Model Cache ────────────────────────────────────────────────────────────────
_cache: dict = {}
_keys_lock = threading.Lock()
_failed_attempts: dict = {}
_fail_lock = threading.Lock()
MAX_FAILS = 10; FAIL_WINDOW = 300

def _check_brute_force(ip):
    now = utcnow().timestamp()
    with _fail_lock:
        a = _failed_attempts.setdefault(ip, deque())
        while a and a[0] < now - FAIL_WINDOW: a.popleft()
        return len(a) >= MAX_FAILS

def _record_fail(ip):
    with _fail_lock:
        _failed_attempts.setdefault(ip, deque()).append(utcnow().timestamp())

def load_u2net():
    if "u2net" in _cache: return _cache["u2net"]
    try:
        from rembg import new_session
        _cache["u2net"] = new_session("u2net")
        logger.info("✅ U2-Net ready"); return _cache["u2net"]
    except Exception as e:
        logger.error(f"U2-Net: {e}"); return None

def load_birefnet():
    if "birefnet" in _cache: return _cache["birefnet"]
    try:
        from transformers import AutoModelForImageSegmentation
        from torchvision import transforms
        m = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
        m = m.float(); m.eval()
        _cache["birefnet"] = m
        _cache["birefnet_tf"] = transforms.Compose([
            transforms.Resize((1024,1024)), transforms.ToTensor(),
            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
        logger.info("✅ BiRefNet-lite ready"); return m
    except Exception as e:
        logger.error(f"BiRefNet: {e}"); return None

import warnings as _warnings
def load_ben2():
    if "ben2" in _cache: return _cache["ben2"]
    try:
        from ben2 import BEN_Base
        with _warnings.catch_warnings():
            _warnings.filterwarnings("ignore", message=".*cuda.*CUDA is not available.*")
            m = BEN_Base.from_pretrained("PramaLLC/BEN2"); m.eval()
        _cache["ben2"] = m
        logger.info("✅ BEN2 ready"); return m
    except Exception as e:
        logger.error(f"BEN2: {e}"); return None

def load_rmbg():
    if "rmbg" in _cache: return _cache["rmbg"]
    try:
        from transformers import AutoModelForImageSegmentation
        from torchvision import transforms
        m = AutoModelForImageSegmentation.from_pretrained(
            "briaai/RMBG-2.0", trust_remote_code=True, token=HF_TOKEN or None)
        m = m.float(); m.eval()
        _cache["rmbg"] = m
        _cache["rmbg_tf"] = transforms.Compose([
            transforms.Resize((1024,1024)), transforms.ToTensor(),
            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
        logger.info("✅ RMBG-2.0 ready"); return m
    except Exception as e:
        logger.error(f"RMBG: {e}"); return None

def _seg(key, img):
    import torch
    m = _cache[key]; tf = _cache[key+"_tf"]; orig = img.size
    inp = tf(img.convert("RGB")).unsqueeze(0).float()
    with torch.no_grad():
        out = m(inp)
        pred = (out[-1] if isinstance(out,(list,tuple)) else out).sigmoid()
    mask = Image.fromarray((pred[0].squeeze().cpu().numpy()*255).astype(np.uint8)).resize(orig, Image.LANCZOS)
    r = img.convert("RGBA"); r.putalpha(mask); return r

def infer_fast(img):
    from rembg import remove
    if not load_u2net(): raise RuntimeError("U2-Net not available")
    buf = io.BytesIO(); img.save(buf,"PNG"); buf.seek(0)
    return Image.open(io.BytesIO(remove(buf.read(), session=_cache["u2net"]))).convert("RGBA")

def infer_quality(img):
    if not load_birefnet(): raise RuntimeError("BiRefNet not available")
    return _seg("birefnet", img)

def infer_best(img):
    if not load_rmbg(): raise RuntimeError("RMBG-2.0 not available")
    return _seg("rmbg", img)

def infer_hair(img):
    if not load_ben2(): raise RuntimeError("BEN2 not available")
    import torch
    _cache["ben2"].to(torch.device("cpu"))
    return _cache["ben2"].inference(img.convert("RGB"), refine_foreground=True)

MODEL_FN   = {"fast": infer_fast, "quality": infer_quality, "best": infer_best, "hair": infer_hair}
MODEL_INFO = {
    "fast":    {"name":"U2-Net",        "desc":"Fastest (~1-5s)"},
    "quality": {"name":"BiRefNet-lite", "desc":"High quality (~5-8s)"},
    "best":    {"name":"BRIA RMBG-2.0", "desc":"Best quality (~20s)"},
    "hair":    {"name":"BEN2",          "desc":"Best for hair (~10-15s)"},
}

# ── Auth ───────────────────────────────────────────────────────────────────────
def validate_key(api_key, model):
    if not api_key: return False, "API key required"
    if api_key not in API_KEYS: return False, "Invalid API key"
    kd = API_KEYS[api_key]; plan = kd.get("plan","free")
    cfg = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
    exp = kd.get("expires_at")
    if exp and utcnow().timestamp() > exp:
        return False, "API key expired. Renew at HDRemover.com"
    if utcnow().timestamp() > kd.get("reset_at", 0):
        kd["calls_today"] = 0
        kd["reset_at"] = (utcnow() + timedelta(days=1)).timestamp()
    now_ts = utcnow().timestamp()
    window = kd.setdefault("minute_window", deque())
    while window and window[0] < now_ts - 60: window.popleft()
    if model not in cfg["models"]:
        return False, f"Model '{model}' not in {plan} plan. Upgrade at HDRemover.com"
    if kd.get("calls_today", 0) >= cfg["daily"]:
        return False, f"Daily limit {cfg['daily']} reached. Resets in 24h."
    if len(window) >= cfg.get("per_minute", 999):
        retry_in = int(60 - (now_ts - window[0])) + 1
        return False, f"Rate limit: {cfg['per_minute']} req/min. Retry in {retry_in}s."
    return True, ""

def _fire_webhook(url, payload):
    import requests as _req
    def _send():
        try: _req.post(url, json=payload, timeout=5); logger.info(f"🔔 Webhook → {url}")
        except Exception as e: logger.warning(f"Webhook failed: {e}")
    threading.Thread(target=_send, daemon=True).start()

def inc_usage(k):
    global _flush_counter
    if k not in API_KEYS: return
    with _keys_lock:
        kd = API_KEYS[k]
        kd["calls_today"] = kd.get("calls_today", 0) + 1
        kd.setdefault("minute_window", deque()).append(utcnow().timestamp())
        _flush_counter += 1
    plan  = kd.get("plan","free"); daily = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])["daily"]
    webhook = kd.get("webhook_url","")
    if webhook and kd["calls_today"] == int(daily * 0.8):
        _fire_webhook(webhook, {"event":"usage_alert_80pct","owner":kd.get("owner",""),
            "plan":plan,"calls_today":kd["calls_today"],"daily_limit":daily,
            "key_hint":f"...{k[-6:]}","timestamp":utcnow().isoformat()})
    _async_save(force=False)

def get_usage(k):
    if k not in API_KEYS: return {"error":"Invalid key"}
    kd = API_KEYS[k]; plan = kd.get("plan","free"); cfg = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
    return {"plan":plan,"calls_today":kd.get("calls_today",0),
            "daily_limit":cfg["daily"],"available_models":cfg["models"],"owner":kd.get("owner","")}

def run_removal(api_key, img_data, model="fast", max_size=0):
    t0 = time.time()
    ok, err = validate_key(api_key, model)
    if not ok: return None, {}, err
    try:
        img = (Image.fromarray(img_data) if isinstance(img_data, np.ndarray)
               else img_data if isinstance(img_data, Image.Image)
               else Image.open(img_data)).convert("RGB")
        ow, oh = img.size
        effective_max = max_size if max_size > 0 else MAX_IMAGE_PX
        if max(ow, oh) > effective_max:
            r = effective_max / max(ow, oh)
            img = img.resize((int(ow*r), int(oh*r)), Image.LANCZOS)
        result = MODEL_FN[model](img)
        inc_usage(api_key)
        elapsed = round(time.time()-t0, 2)
        meta = {"model":MODEL_INFO[model]["name"],"model_key":model,
                "original_size":f"{ow}x{oh}","output_size":f"{result.size[0]}x{result.size[1]}",
                "processing_time_sec":elapsed,"calls_today":API_KEYS[api_key]["calls_today"]}
        logger.info(f"✅ {model} {ow}x{oh} {elapsed}s key=...{api_key[-6:]}")
        return result, meta, ""
    except Exception as e:
        logger.error(f"run_removal: {e}", exc_info=True)
        return None, {}, str(e)

# ── Admin ──────────────────────────────────────────────────────────────────────
import gradio as gr

_CHOICES = ["🚀 Fast (U2-Net) ~1-5s","⚡ Quality (BiRefNet-lite) ~5-8s",
            "💇 Hair & Portrait (BEN2) ~10-15s","🏆 Best Quality (RMBG-2.0) ~20s"]
_CMAP    = {"🚀 Fast (U2-Net) ~1-5s":"fast","⚡ Quality (BiRefNet-lite) ~5-8s":"quality",
            "💇 Hair & Portrait (BEN2) ~10-15s":"hair","🏆 Best Quality (RMBG-2.0) ~20s":"best"}

def ui_usage(api_key):
    info = get_usage(api_key)
    if "error" in info: return f"❌ {info['error']}"
    return (f"**Plan:** {info['plan'].upper()}  \n**Calls today:** {info['calls_today']} / {info['daily_limit']}  \n"
            f"**Models:** {', '.join(info['available_models'])}  \n**Owner:** {info.get('owner','—')}")

def admin_login(password):
    if password == ADMIN_PASS:
        return gr.update(visible=False), gr.update(visible=True), "", admin_list_customers(), admin_stats()
    return gr.update(visible=True), gr.update(visible=False), "❌ Wrong password", [], ""

def admin_list_customers():
    rows = []
    for key, d in API_KEYS.items():
        if key == MASTER_KEY: continue
        exp = d.get("expires_at")
        exp_str = datetime.fromtimestamp(exp, tz=timezone.utc).strftime("%Y-%m-%d") if exp else "Never"
        rows.append([key, d.get("owner",""), d.get("plan","free"),
                     d.get("calls_today",0), d.get("limit", PLAN_LIMITS.get(d.get("plan","free"),{}).get("daily",10)),
                     d.get("created_at","")[:10] if d.get("created_at") else "", exp_str,
                     "✅" if d.get("webhook_url") else "—"])
    return rows

def admin_add_customer(email, plan, custom_limit, expiry_days="", webhook_url=""):
    if not email or "@" not in email: return "❌ Valid email required", admin_list_customers(), ""
    if plan not in PLAN_LIMITS: return "❌ Invalid plan", admin_list_customers(), ""
    key = gen_key(plan)
    limit = int(custom_limit) if str(custom_limit).strip().isdigit() else PLAN_LIMITS[plan]["daily"]
    days = int(expiry_days) if str(expiry_days).strip().isdigit() and int(expiry_days) > 0 else PLAN_DEFAULT_EXPIRY.get(plan)
    expires_at = (utcnow() + timedelta(days=days)).timestamp() if days else None
    API_KEYS[key] = {"plan":plan,"owner":email,"calls_today":0,"reset_at":0,"limit":limit,
                     "models":PLAN_LIMITS[plan]["models"],"created_at":utcnow().isoformat(),
                     "expires_at":expires_at,"webhook_url":webhook_url.strip() if webhook_url else ""}
    ok, msg = push_keys_to_hf_secret()
    return (f"✅ Key for **{email}** ({plan})\n\n🔑 `{key}`\n⏳ Expires: **{f'{days} days' if days else 'Never'}**\n\n{msg}",
            admin_list_customers(), key)

def admin_delete_customer(key):
    if not key or key.strip() not in API_KEYS: return "❌ Key not found", admin_list_customers()
    key = key.strip()
    if key == MASTER_KEY: return "❌ Cannot delete master key", admin_list_customers()
    owner = API_KEYS[key].get("owner",""); del API_KEYS[key]
    ok, msg = push_keys_to_hf_secret()
    return f"✅ Deleted **{owner}**\n\n{msg}", admin_list_customers()

def admin_upgrade_plan(key, new_plan):
    if not key or key.strip() not in API_KEYS: return "❌ Key not found", admin_list_customers()
    if new_plan not in PLAN_LIMITS: return "❌ Invalid plan", admin_list_customers()
    key = key.strip()
    API_KEYS[key]["plan"] = new_plan; API_KEYS[key]["limit"] = PLAN_LIMITS[new_plan]["daily"]
    API_KEYS[key]["models"] = PLAN_LIMITS[new_plan]["models"]
    ok, msg = push_keys_to_hf_secret()
    return f"✅ **{API_KEYS[key].get('owner','')}** → **{new_plan}**\n\n{msg}", admin_list_customers()

def admin_stats():
    vals = list(API_KEYS.values())
    total = len([v for v in vals if v.get("owner") != "HDRemover.com"])
    free_c = len([v for v in vals if v.get("plan")=="free"])
    start  = len([v for v in vals if v.get("plan")=="starter"])
    pro_c  = len([v for v in vals if v.get("plan") in ("pro","master") and v.get("owner")!="HDRemover.com"])
    calls  = sum(v.get("calls_today",0) for v in vals)
    return (f"👥 **Total:** {total}  \n🆓 **Free:** {free_c}  \n⭐ **Starter:** {start}  \n"
            f"🏆 **Pro:** {pro_c}  \n📊 **Calls Today:** {calls}  \n💰 **Est. MRR:** ${start*9+pro_c*29}")

def admin_export_json():
    safe = {}
    for k, v in API_KEYS.items():
        entry = dict(v)
        if "minute_window" in entry: entry["minute_window"] = list(entry["minute_window"])
        safe[k] = entry
    return json.dumps(safe, indent=2, default=str)

# ── BG Helper ──────────────────────────────────────────────────────────────────
import zipfile, tempfile

_session_history: dict = {}

def _history_add(api_key, entry):
    if api_key not in _session_history: _session_history[api_key] = []
    _session_history[api_key].append(entry)
    if len(_session_history[api_key]) > 50: _session_history[api_key].pop(0)

def _history_get(api_key): return _session_history.get(api_key, [])
def _history_clear(api_key): _session_history[api_key] = []

def _apply_background(result_rgba, bg_choice, bg_color, bg_image=None):
    if bg_choice == "transparent": return result_rgba
    w, h = result_rgba.size
    if bg_choice == "color":
        try:
            from PIL import ImageColor; rgb = ImageColor.getrgb(bg_color)
        except: rgb = (255,255,255)
        canvas = Image.new("RGBA",(w,h),rgb+(255,)); canvas.paste(result_rgba,mask=result_rgba.split()[3])
        return canvas.convert("RGB")
    if bg_choice == "custom_image" and bg_image is not None:
        bg = (Image.fromarray(bg_image) if isinstance(bg_image,np.ndarray) else bg_image).convert("RGBA")
        bg = bg.resize((w,h),Image.LANCZOS); bg.paste(result_rgba,mask=result_rgba.split()[3])
        return bg.convert("RGB")
    return result_rgba

def ui_remove_with_bg(api_key, image, choice, bg_choice, bg_color, bg_image):
    m = _CMAP.get(choice,"fast")
    result, meta, err = run_removal(api_key, image, m)
    if err: return None, f"❌ **Error:** {err}"
    final = _apply_background(result, bg_choice, bg_color, bg_image)
    _history_add(api_key or "anon", {"name":"single_image","time":utcnow().strftime("%H:%M:%S"),
                                      "model":meta["model"],"size":meta["output_size"],"result":final})
    return final, (f"✅ **Done** | `{meta['model']}` | `{meta['processing_time_sec']}s` | "
                   f"`{meta['original_size']}` → `{meta['output_size']}` | calls: `{meta['calls_today']}`")

def ui_batch_remove(api_key, images, choice, bg_choice, bg_color, bg_image, progress=gr.Progress()):
    if not images: return [], "⚠️ No images uploaded.", None
    model = _CMAP.get(choice,"fast"); results_gallery=[]; status_lines=[]; processed_pils=[]
    progress(0, desc="Starting batch…")
    for i, img_input in enumerate(images):
        fname = f"image_{i+1}"
        try:
            if isinstance(img_input, str): pil_in = Image.open(img_input).convert("RGB"); fname = os.path.basename(img_input)
            elif isinstance(img_input, np.ndarray): pil_in = Image.fromarray(img_input).convert("RGB")
            elif isinstance(img_input, Image.Image): pil_in = img_input.convert("RGB")
            else: status_lines.append(f"❌ `{fname}` — unsupported"); continue
        except Exception as e: status_lines.append(f"❌ `{fname}` — {e}"); progress((i+1)/len(images)); continue
        progress(i/len(images), desc=f"Processing {i+1}/{len(images)}: {fname}")
        result, meta, err = run_removal(api_key, pil_in, model)
        if err: status_lines.append(f"❌ `{fname}` — {err}"); progress((i+1)/len(images)); continue
        final = _apply_background(result, bg_choice, bg_color, bg_image)
        _history_add(api_key or "anon", {"name":fname,"time":utcnow().strftime("%H:%M:%S"),
                                          "model":meta["model"],"size":meta["output_size"],"result":final})
        results_gallery.append(final); processed_pils.append((fname,final))
        status_lines.append(f"✅ `{fname}` — {meta['processing_time_sec']}s | {meta['output_size']}")
        progress((i+1)/len(images), desc=f"Done: {fname}")
    zip_path = None
    if processed_pils:
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
        with zipfile.ZipFile(tmp.name,"w",zipfile.ZIP_DEFLATED) as zf:
            for fname, pil in processed_pils:
                buf = io.BytesIO(); ext = "png" if bg_choice=="transparent" else "jpg"
                pil.save(buf,"PNG" if bg_choice=="transparent" else "JPEG",optimize=True); buf.seek(0)
                zf.writestr(os.path.splitext(fname)[0]+f"_nobg.{ext}", buf.read())
        zip_path = tmp.name
    total = len(processed_pils); failed = len(images)-total
    return results_gallery, f"### Batch Complete\n✅ **{total}** processed | ❌ **{failed}** failed\n\n"+"\n".join(status_lines), zip_path

def ui_get_history(api_key):
    items = _history_get(api_key or "anon")
    if not items: return [], "No images processed yet."
    gallery = [item["result"] for item in items]
    lines = [f"| `{it['name']}` | {it['model']} | {it['size']} | {it['time']} |" for it in items]
    return gallery, f"### History ({len(items)})\n\n| File | Model | Size | Time |\n|---|---|---|---|\n"+"\n".join(lines)

def ui_clear_history(api_key):
    _history_clear(api_key or "anon"); return [], "🗑️ Cleared."

# ── Gradio UI ──────────────────────────────────────────────────────────────────
with gr.Blocks(title="HDRemover — Background Removal API") as demo:
    gr.HTML("""<div style="text-align:center;padding:18px 0 8px">
      <h1 style="font-size:26px;font-weight:900;margin:0">🎨 HDRemover — Background Removal API</h1>
      <p style="color:#888;font-size:13px;margin:6px 0 0">4 AI Models · API Key Auth ·
        <a href="https://hdremover.com" target="_blank" style="color:#00e5a0">HDRemover.com</a></p>
    </div>""")

    with gr.Tabs():
        with gr.TabItem("🖼️ Remove Background"):
            with gr.Row():
                with gr.Column():
                    t_key   = gr.Textbox(label="🔑 API Key", type="password", placeholder="hdremover-xxxx-xxxx")
                    t_model = gr.Dropdown(label="Model", choices=_CHOICES, value=_CHOICES[0])
                    t_img   = gr.Image(label="Upload Image", type="pil")
                    t_btn   = gr.Button("✨ Remove Background", variant="primary", size="lg")
                with gr.Column():
                    t_out  = gr.Image(label="Result — Transparent PNG", type="pil")
                    t_info = gr.Markdown("*Upload an image and click Remove Background*")
            with gr.Accordion("🎨 Background Options (optional)", open=False):
                t_bg_choice = gr.Radio(label="Output Background", choices=["transparent","color","custom_image"], value="transparent")
                t_bg_color  = gr.ColorPicker(label="Color", value="#ffffff", visible=False)
                t_bg_image  = gr.Image(label="Custom Background", type="pil", visible=False)
            def _t_toggle(c): return gr.update(visible=c=="color"), gr.update(visible=c=="custom_image")
            t_bg_choice.change(_t_toggle, [t_bg_choice], [t_bg_color, t_bg_image])
            t_btn.click(ui_remove_with_bg, [t_key,t_img,t_model,t_bg_choice,t_bg_color,t_bg_image], [t_out,t_info])
            gr.Markdown("---\n### Plans\n| Plan | Daily | Models | Price |\n|---|---|---|---|\n| Free | 10/day | fast | $0 |\n| Starter | 100/day | fast+quality+hair | $9/mo |\n| Pro | 500/day | All 4 | $29/mo |\n\n📧 **support@hdremover.com**")

        with gr.TabItem("📊 My Usage"):
            u_key = gr.Textbox(label="🔑 Your API Key", type="password")
            u_btn = gr.Button("Check Usage", variant="secondary")
            u_out = gr.Markdown()
            u_btn.click(ui_usage, [u_key], [u_out])

        with gr.TabItem("📚 API Docs"):
            gr.Markdown("""## REST API\n**Base URL:** `https://hdremover-background-remover.hf.space`\n\n### POST /api/remove-bg\n```\nHeader: X-API-Key: your-key\nBody:   image=<file>, model=fast|quality|hair|best\n```\n\n### GET /api/usage\n```\nHeader: X-API-Key: your-key\n```\n\n### GET /health\nReturns server status and loaded models.\n\n### cURL Example\n```bash\ncurl -X POST https://hdremover-background-remover.hf.space/api/remove-bg \\\n  -H "X-API-Key: your-key" \\\n  -F "image=@photo.jpg" \\\n  -F "model=fast" \\\n  --output result.png\n```""")

        with gr.TabItem("🔐 Admin"):
            with gr.Column(visible=True) as admin_login_col:
                gr.Markdown("### 🔐 Admin Login")
                a_pwd  = gr.Textbox(label="Password", type="password")
                a_lbtn = gr.Button("Login", variant="primary")
                a_lerr = gr.Markdown()
            with gr.Column(visible=False) as admin_panel_col:
                gr.Markdown("## 👑 Admin Panel")
                with gr.Tabs():
                    with gr.TabItem("➕ Add Customer"):
                        with gr.Row():
                            a_email = gr.Textbox(label="Email", scale=2)
                            a_plan  = gr.Dropdown(label="Plan", choices=["free","starter","pro","master"], value="starter", scale=1)
                            a_limit = gr.Textbox(label="Custom Limit (optional)", scale=1)
                        with gr.Row():
                            a_expiry  = gr.Textbox(label="Expires In (days)", scale=1)
                            a_webhook = gr.Textbox(label="Webhook URL (optional)", scale=3)
                        a_add_btn = gr.Button("✨ Generate Key + Save", variant="primary")
                        a_add_out = gr.Markdown()
                        a_new_key = gr.Textbox(label="🔑 Generated Key", interactive=True)
                    with gr.TabItem("👥 Customers"):
                        a_refresh_btn = gr.Button("🔄 Refresh", size="sm")
                        a_table = gr.Dataframe(
                            headers=["API Key","Owner","Plan","Calls","Limit","Created","Expires","Webhook"],
                            datatype=["str","str","str","number","number","str","str","str"], interactive=False, wrap=True)
                        with gr.Row():
                            a_edit_key  = gr.Textbox(label="Key to Edit")
                            a_edit_plan = gr.Dropdown(label="New Plan", choices=["free","starter","pro","master"], value="pro")
                        a_upgrade_btn = gr.Button("⬆️ Upgrade Plan", variant="secondary")
                        a_upgrade_out = gr.Markdown()
                        with gr.Row():
                            a_del_key = gr.Textbox(label="Key to Delete")
                            a_del_btn = gr.Button("🗑️ Delete", variant="stop")
                        a_del_out = gr.Markdown()
                    with gr.TabItem("📊 Stats"):
                        a_stats_btn = gr.Button("🔄 Refresh Stats")
                        a_stats_out = gr.Markdown()
                        a_stats_btn.click(admin_stats, [], [a_stats_out])
                    with gr.TabItem("📋 Export JSON"):
                        a_export_btn = gr.Button("📋 Export")
                        a_export_out = gr.Code(language="json", interactive=False)
                        a_export_btn.click(admin_export_json, [], [a_export_out])

            a_lbtn.click(admin_login, [a_pwd], [admin_login_col,admin_panel_col,a_lerr,a_table,a_stats_out])
            a_pwd.submit(admin_login, [a_pwd], [admin_login_col,admin_panel_col,a_lerr,a_table,a_stats_out])
            a_add_btn.click(admin_add_customer, [a_email,a_plan,a_limit,a_expiry,a_webhook], [a_add_out,a_table,a_new_key])
            a_refresh_btn.click(lambda: admin_list_customers(), [], [a_table])
            a_upgrade_btn.click(admin_upgrade_plan, [a_edit_key,a_edit_plan], [a_upgrade_out,a_table])
            a_del_btn.click(admin_delete_customer, [a_del_key], [a_del_out,a_table])

        with gr.TabItem("📦 Batch Processing"):
            with gr.Row():
                with gr.Column():
                    b_key   = gr.Textbox(label="🔑 API Key", type="password")
                    b_model = gr.Dropdown(label="Model", choices=_CHOICES, value=_CHOICES[0])
                    b_files = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
                    b_bg_choice = gr.Radio(label="Output BG", choices=["transparent","color","custom_image"], value="transparent")
                    b_bg_color  = gr.ColorPicker(label="Color", value="#ffffff", visible=False)
                    b_bg_image  = gr.Image(label="Custom BG", type="pil", visible=False)
                    b_btn = gr.Button("🚀 Start Batch", variant="primary", size="lg")
                with gr.Column():
                    b_gallery = gr.Gallery(label="Results", columns=3, object_fit="contain", height=360)
                    b_status  = gr.Markdown("*Upload images and click Start Batch*")
                    b_zip     = gr.File(label="⬇️ Download ZIP")
            def _tb(c): return gr.update(visible=c=="color"), gr.update(visible=c=="custom_image")
            b_bg_choice.change(_tb, [b_bg_choice], [b_bg_color,b_bg_image])
            b_btn.click(ui_batch_remove, [b_key,b_files,b_model,b_bg_choice,b_bg_color,b_bg_image], [b_gallery,b_status,b_zip])

        with gr.TabItem("🎨 BG Replace"):
            with gr.Row():
                with gr.Column():
                    br_key   = gr.Textbox(label="🔑 API Key", type="password")
                    br_model = gr.Dropdown(label="Model", choices=_CHOICES, value=_CHOICES[0])
                    br_img   = gr.Image(label="Upload Image", type="pil")
                    br_bg_choice = gr.Radio(label="New Background", choices=["transparent","color","custom_image"], value="color")
                    br_bg_color  = gr.ColorPicker(label="Color", value="#ffffff")
                    br_bg_image  = gr.Image(label="Custom BG Image", type="pil", visible=False)
                    br_btn = gr.Button("✨ Remove & Replace", variant="primary", size="lg")
                with gr.Column():
                    br_out  = gr.Image(label="Result", type="pil")
                    br_info = gr.Markdown("*Upload an image and choose a background*")
            def _brb(c): return gr.update(visible=c=="color"), gr.update(visible=c=="custom_image")
            br_bg_choice.change(_brb, [br_bg_choice], [br_bg_color,br_bg_image])
            br_btn.click(ui_remove_with_bg, [br_key,br_img,br_model,br_bg_choice,br_bg_color,br_bg_image], [br_out,br_info])

        with gr.TabItem("🕒 History"):
            with gr.Row():
                h_refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
                h_clear_btn   = gr.Button("🗑️ Clear", variant="stop")
            h_gallery = gr.Gallery(label="Processed Images", columns=4, object_fit="contain", height=400)
            h_status  = gr.Markdown("*Click Refresh to load history*")
            h_key = gr.Textbox(label="🔑 API Key", type="password")
            h_refresh_btn.click(ui_get_history, [h_key], [h_gallery,h_status])
            h_clear_btn.click(ui_clear_history, [h_key], [h_gallery,h_status])

# ── FastAPI Routes ─────────────────────────────────────────────────────────────
from fastapi import Request
from fastapi.responses import Response, JSONResponse
from fastapi.middleware.cors import CORSMiddleware

demo.app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET","POST","OPTIONS"],
    allow_headers=["*"],
    expose_headers=["X-Model","X-Processing-Time","X-Original-Size","X-Output-Size","X-Calls-Today"],
)

@demo.app.post("/api/remove-bg")
async def api_remove_bg(request: Request):
    client_ip = request.headers.get("x-forwarded-for", request.client.host or "unknown").split(",")[0].strip()
    if _check_brute_force(client_ip):
        return JSONResponse({"success":False,"error":"Too many failed attempts."}, status_code=429)
    x_api_key = request.headers.get("x-api-key") or request.headers.get("X-API-Key")
    if not x_api_key:
        _record_fail(client_ip)
        return JSONResponse({"success":False,"error":"X-API-Key header required"}, status_code=401)
    try:
        form = await request.form()
        model = str(form.get("model","fast")); max_size = int(form.get("max_size",0))
        img_file = form.get("image")
        if not img_file: return JSONResponse({"success":False,"error":"'image' field required"}, status_code=400)
        raw = await img_file.read()
        if len(raw) > MAX_IMAGE_BYTES:
            return JSONResponse({"success":False,"error":f"File too large. Max {MAX_IMAGE_BYTES//(1024*1024)}MB."}, status_code=413)
        try: pil = Image.open(io.BytesIO(raw)).convert("RGB")
        except: return JSONResponse({"success":False,"error":"Invalid image file."}, status_code=400)
    except Exception as e:
        return JSONResponse({"success":False,"error":f"Bad request: {e}"}, status_code=400)
    if model not in MODEL_FN:
        return JSONResponse({"success":False,"error":f"Invalid model. Choose: {list(MODEL_FN)}"}, status_code=400)
    result, meta, err = run_removal(x_api_key, pil, model, max_size)
    if err:
        is_auth = any(w in err.lower() for w in ("key","limit","plan","upgrade","invalid","expired"))
        if is_auth: _record_fail(client_ip)
        return JSONResponse({"success":False,"error":err}, status_code=401 if is_auth else 500)
    buf = io.BytesIO(); result.save(buf,"PNG",optimize=True); buf.seek(0)
    return Response(buf.read(), media_type="image/png", headers={
        "X-Model":meta["model"],"X-Processing-Time":str(meta["processing_time_sec"]),
        "X-Original-Size":meta["original_size"],"X-Output-Size":meta["output_size"],
        "X-Calls-Today":str(meta["calls_today"]),
    })

@demo.app.get("/api/usage")
async def api_usage(request: Request):
    key = request.headers.get("x-api-key") or request.headers.get("X-API-Key")
    if not key: return JSONResponse({"error":"X-API-Key required"}, status_code=401)
    info = get_usage(key)
    return JSONResponse({"success":"error" not in info, **info}, status_code=200 if "error" not in info else 401)

@demo.app.get("/api/models")
async def api_models():
    return {"success":True,"models":{k:v for k,v in MODEL_INFO.items()}}

@demo.app.get("/api/health")
async def health():
    return {"status":"ok","hf_token":"set" if HF_TOKEN else "missing",
            "models_ready":[k for k in _cache if not k.endswith("_tf")],
            "timestamp":utcnow().isoformat()}

# ── Launch ─────────────────────────────────────────────────────────────────────
def _periodic_flush():
    import time as _t
    while True:
        _t.sleep(FLUSH_EVERY_S)
        save_keys_to_dataset(force=True)

threading.Thread(target=_periodic_flush, daemon=True).start()
logger.info("⏰ Periodic flush started")
logger.info("Preloading U2-Net (blocking)...")
load_u2net()

def _bg_preload():
    logger.info("Background preload: BiRefNet-lite...")
    load_birefnet()
    logger.info("Background preload: BEN2...")
    load_ben2()
    logger.info("Background preload: RMBG-2.0...")
    load_rmbg()
    logger.info("✅ All 4 models ready — zero cold starts")

threading.Thread(target=_bg_preload, daemon=True).start()

demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False, strict_cors=False)