"""
FreeMG — Magic Grab API
========================
Space: freebg/magic-grab
CPU Mode  : GrabCut segmentation + cv2 inpainting (works on free HF CPU space)
GPU Mode  : SAM2.1-hiera-large + SDXL inpainting (auto-enabled if GPU available)
Auth      : Master key (unlimited) + per-customer keys with daily quotas
Admin     : Full dashboard — add/delete/upgrade customers, stats, export
Auto-save : Customer keys persisted to HF_SECRET via huggingface_hub
REST      : POST /api/magic-grab  (mounted on Gradio FastAPI app)
"""

# ─────────────────────────────────────────────────────────────────────────────
# Imports
# ─────────────────────────────────────────────────────────────────────────────
import gradio as gr
import numpy as np
import torch
import os, json, time, uuid, base64, traceback
from PIL import Image
from datetime import datetime, date
from io import BytesIO
from fastapi import File, UploadFile, Form
from fastapi.responses import JSONResponse

# ── Detect GPU availability ───────────────────────────────────────────────────
HAS_GPU = torch.cuda.is_available()

# ── spaces decorator — graceful fallback when no ZeroGPU ─────────────────────
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False
    class spaces:                          # noqa: F811
        @staticmethod
        def GPU(duration=60):
            def decorator(fn): return fn
            return decorator

print(f"[FreeMG] GPU available: {HAS_GPU} | spaces: {HAS_SPACES}")

# ─────────────────────────────────────────────────────────────────────────────
# Environment / Secrets
# ─────────────────────────────────────────────────────────────────────────────
MASTER_API_KEY = os.environ.get("MASTER_API_KEY", "")
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "admin123")
HF_TOKEN       = os.environ.get("HF_TOKEN", "")
HF_REPO_ID     = os.environ.get("HF_REPO_ID", "freebg/magic-grab")

def _load_api_keys() -> dict:
    raw = os.environ.get("API_KEYS_JSON", "{}")
    try:
        return json.loads(raw)
    except Exception:
        return {}

API_KEYS: dict = _load_api_keys()

PLAN_LIMITS = {"free": 10, "starter": 100, "pro": 500, "master": 999_999}
PLAN_PRICES = {"free":  0, "starter":   9, "pro":  29, "master":       0}

# ─────────────────────────────────────────────────────────────────────────────
# Usage tracker  (in-memory; resets on Space restart — once per day is fine)
# ─────────────────────────────────────────────────────────────────────────────
_usage_today: dict = {}
_usage_date:  str  = str(date.today())

def _reset_if_new_day():
    global _usage_today, _usage_date
    today = str(date.today())
    if today != _usage_date:
        _usage_today = {}
        _usage_date  = today

def _used_today(key: str) -> int:
    _reset_if_new_day()
    return _usage_today.get(key, 0)

def _record_usage(key: str):
    _reset_if_new_day()
    _usage_today[key] = _usage_today.get(key, 0) + 1

# ─────────────────────────────────────────────────────────────────────────────
# API Key validation
# ─────────────────────────────────────────────────────────────────────────────
def validate_key(api_key: str) -> tuple:
    """Returns (is_valid: bool, reason: str, info: dict)"""
    api_key = (api_key or "").strip()
    if not api_key:
        return False, "no_key", {}
    if api_key == MASTER_API_KEY and MASTER_API_KEY:
        return True, "master", {"email": "master@freebg.site", "plan": "master", "limit": 999_999}
    if api_key in API_KEYS:
        info  = API_KEYS[api_key]
        plan  = info.get("plan", "free")
        limit = info.get("custom_limit") or PLAN_LIMITS.get(plan, 10)
        used  = _used_today(api_key)
        if used >= limit:
            return False, "quota_exceeded", info
        return True, "ok", info
    return False, "invalid_key", {}

# ─────────────────────────────────────────────────────────────────────────────
# HF Secret auto-save
# ─────────────────────────────────────────────────────────────────────────────
def _save_hf_secret(name: str, value: str) -> bool:
    if not HF_TOKEN or not HF_REPO_ID:
        print(f"[HF] No token/repo_id — skipping secret save for {name}")
        return False
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        api.add_space_secret(repo_id=HF_REPO_ID, key=name, value=value)
        print(f"[HF] ✓ Secret '{name}' saved to {HF_REPO_ID}")
        return True
    except Exception as e:
        print(f"[HF] ✗ Failed to save secret '{name}': {e}")
        return False

def _persist_keys() -> bool:
    return _save_hf_secret("API_KEYS_JSON", json.dumps(API_KEYS))

# ─────────────────────────────────────────────────────────────────────────────
# Model singletons  (GPU only — not loaded on CPU space)
# ─────────────────────────────────────────────────────────────────────────────
_sam2_processor = None
_sam2_model     = None
_sdxl_pipe      = None

def _load_sam2():
    global _sam2_processor, _sam2_model
    if _sam2_model is not None:
        return
    from transformers import Sam2Processor, Sam2Model
    print("[SAM2] Loading facebook/sam2.1-hiera-large …")
    _sam2_processor = Sam2Processor.from_pretrained("facebook/sam2.1-hiera-large")
    _sam2_model = Sam2Model.from_pretrained(
        "facebook/sam2.1-hiera-large",
        torch_dtype=torch.float16,
    ).to("cuda")
    _sam2_model.eval()
    print("[SAM2] ✓ Ready")

def _load_sdxl():
    global _sdxl_pipe
    if _sdxl_pipe is not None:
        return
    from diffusers import AutoPipelineForInpainting
    print("[SDXL] Loading inpainting pipeline …")
    _sdxl_pipe = AutoPipelineForInpainting.from_pretrained(
        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
        torch_dtype=torch.float16,
        variant="fp16",
    ).to("cuda")
    try:
        _sdxl_pipe.enable_xformers_memory_efficient_attention()
    except Exception:
        pass
    print("[SDXL] ✓ Ready")

# ─────────────────────────────────────────────────────────────────────────────
# Segmentation — CPU: GrabCut  |  GPU: SAM2.1
# ─────────────────────────────────────────────────────────────────────────────
def _segment_grabcut(image: Image.Image, points: list) -> Image.Image:
    """
    CPU segmentation via OpenCV GrabCut.
    points: [[x, y], ...]  — first point is treated as subject center.
    Builds a bounding rect around the click cluster (or uses center ±30% if 1 point).
    Returns binary mask (PIL 'L', 0/255).
    """
    import cv2
    img_np  = np.array(image.convert("RGB"))
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    H, W    = img_bgr.shape[:2]

    # Build rect from points
    if len(points) == 1:
        cx, cy = points[0]
        pad_x  = int(W * 0.30)
        pad_y  = int(H * 0.30)
        x1 = max(0,  cx - pad_x)
        y1 = max(0,  cy - pad_y)
        x2 = min(W,  cx + pad_x)
        y2 = min(H,  cy + pad_y)
    else:
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        pad = 30
        x1 = max(0, min(xs) - pad)
        y1 = max(0, min(ys) - pad)
        x2 = min(W, max(xs) + pad)
        y2 = min(H, max(ys) + pad)

    rect  = (x1, y1, x2 - x1, y2 - y1)
    mask  = np.zeros(img_bgr.shape[:2], np.uint8)
    bgd_m = np.zeros((1, 65), np.float64)
    fgd_m = np.zeros((1, 65), np.float64)
    cv2.grabCut(img_bgr, mask, rect, bgd_m, fgd_m, 5, cv2.GC_INIT_WITH_RECT)

    binary = np.where(
        (mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0
    ).astype(np.uint8)
    return Image.fromarray(binary, mode="L")

@spaces.GPU(duration=30)
def _segment_sam2(image: Image.Image, points: list, labels: list) -> Image.Image:
    """GPU segmentation via SAM2.1-hiera-large. Returns PIL 'L' mask."""
    _load_sam2()
    inputs = _sam2_processor(
        images=image,
        input_points=[points],
        input_labels=[labels],
        return_tensors="pt",
    ).to("cuda", torch.float16)

    with torch.no_grad():
        outputs = _sam2_model(**inputs)

    masks = _sam2_processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu(),
    )
    scores   = outputs.iou_scores[0, 0].cpu().numpy()
    best_idx = int(np.argmax(scores))
    mask_np  = masks[0][0][best_idx].numpy().astype(np.uint8) * 255
    return Image.fromarray(mask_np, mode="L")

def _run_segmentation(image: Image.Image, points: list, labels: list) -> Image.Image:
    """Router: SAM2 on GPU, GrabCut on CPU."""
    if HAS_GPU:
        return _segment_sam2(image, points, labels)
    return _segment_grabcut(image, points)

# ─────────────────────────────────────────────────────────────────────────────
# Inpainting — CPU: cv2 Telea  |  GPU: SDXL
# ─────────────────────────────────────────────────────────────────────────────
def _inpaint_cpu(image: Image.Image, mask: Image.Image) -> Image.Image:
    """Fast CPU inpainting using OpenCV Telea algorithm."""
    import cv2
    img_np  = np.array(image.convert("RGB"))
    mask_np = np.array(mask.convert("L"))
    kernel      = np.ones((15, 15), np.uint8)
    dilated     = cv2.dilate(mask_np, kernel, iterations=2)
    result      = cv2.inpaint(img_np, dilated, inpaintRadius=5, flags=cv2.INPAINT_TELEA)
    return Image.fromarray(result)

@spaces.GPU(duration=60)
def _inpaint_sdxl(image: Image.Image, mask: Image.Image, prompt: str, steps: int) -> Image.Image:
    """SDXL inpainting — quality / best mode."""
    _load_sdxl()
    W, H         = image.size
    img_resized  = image.resize((1024, 1024), Image.LANCZOS)
    mask_resized = mask.resize((1024, 1024), Image.NEAREST)

    if not prompt.strip():
        prompt = "seamless natural background, photorealistic, ultra HD, no artifacts"
    neg = "blurry, distorted, watermark, text, extra limbs, duplicate"

    result = _sdxl_pipe(
        prompt=prompt,
        negative_prompt=neg,
        image=img_resized,
        mask_image=mask_resized,
        num_inference_steps=steps,
        guidance_scale=7.5,
        strength=0.99,
    ).images[0]
    return result.resize((W, H), Image.LANCZOS)

def _run_inpainting(image, mask, prompt, mode) -> Image.Image:
    """Router: SDXL on GPU, cv2 on CPU. Fast mode always uses cv2."""
    if mode == "fast" or not HAS_GPU:
        return _inpaint_cpu(image, mask)
    steps = 20 if mode == "quality" else 35
    try:
        return _inpaint_sdxl(image, mask, prompt, steps)
    except Exception as e:
        print(f"[Inpaint] SDXL failed ({e}), falling back to CPU")
        return _inpaint_cpu(image, mask)

# ─────────────────────────────────────────────────────────────────────────────
# Subject extraction
# ─────────────────────────────────────────────────────────────────────────────
def _extract_subject(image: Image.Image, mask: Image.Image) -> Image.Image:
    """Cut subject out with transparent background (RGBA PNG)."""
    img_rgba = image.convert("RGBA")
    r, g, b, _ = img_rgba.split()
    return Image.merge("RGBA", (r, g, b, mask.convert("L")))

# ─────────────────────────────────────────────────────────────────────────────
# Utility helpers
# ─────────────────────────────────────────────────────────────────────────────
def _pil_to_b64(img: Image.Image, fmt: str = "PNG") -> str:
    buf = BytesIO()
    img.save(buf, format=fmt)
    return base64.b64encode(buf.getvalue()).decode()

def _parse_image(inp) -> Image.Image:
    if isinstance(inp, np.ndarray):
        return Image.fromarray(inp).convert("RGB")
    if isinstance(inp, str):
        return Image.open(BytesIO(base64.b64decode(inp))).convert("RGB")
    return inp.convert("RGB")

def _parse_points(pts_json: str) -> tuple:
    """Returns (points [[x,y],...], labels [1/0,...]) or (None, None) on failure."""
    try:
        data   = json.loads(pts_json or "[]")
        points = [[int(p["x"]), int(p["y"])] for p in data]
        labels = [int(p.get("label", 1))     for p in data]
        if not points:
            raise ValueError("empty")
        return points, labels
    except Exception:
        return None, None

# ─────────────────────────────────────────────────────────────────────────────
# Main pipeline
# ─────────────────────────────────────────────────────────────────────────────
def magic_grab_pipeline(
    image_input,
    click_points_json: str,
    mode: str,
    fill_prompt: str,
    output_type: str,
    api_key: str = "",
) -> tuple:
    """
    Returns: (subject_img, filled_bg_img, composite_img, status_str)
    """
    t0 = time.time()

    # Auth
    if api_key.strip():
        valid, reason, _ = validate_key(api_key)
        if not valid:
            msg = "❌ Daily quota exceeded. Upgrade your plan." \
                  if reason == "quota_exceeded" \
                  else "❌ Invalid API key. Get one at freebg.site"
            return None, None, None, msg

    # Parse image
    try:
        image = _parse_image(image_input)
    except Exception as e:
        return None, None, None, f"❌ Could not read image: {e}"

    W, H = image.size

    # Parse click points
    points, labels = _parse_points(click_points_json)
    if points is None:
        points = [[W // 2, H // 2]]
        labels = [1]

    # Step 1 — Segment
    try:
        mask = _run_segmentation(image, points, labels)
    except Exception as e:
        return None, None, None, f"❌ Segmentation failed: {e}\n{traceback.format_exc()}"

    # Step 2 — Extract subject
    subject_img = _extract_subject(image, mask)

    # Step 3 — Inpaint background
    try:
        filled_bg = _run_inpainting(image, mask, fill_prompt, mode)
    except Exception as e:
        return None, None, None, f"❌ Inpainting failed: {e}"

    # Step 4 — Composite
    composite = filled_bg.convert("RGBA")
    composite.paste(subject_img, (0, 0), subject_img)
    composite = composite.convert("RGB")

    # Record usage
    if api_key.strip():
        _record_usage(api_key.strip())

    elapsed   = round(time.time() - t0, 2)
    seg_tag   = "SAM2.1" if HAS_GPU else "GrabCut(CPU)"
    fill_tag  = {"fast": "cv2-Telea", "quality": "SDXL-20steps", "best": "SDXL-35steps"}.get(mode, mode) \
                if HAS_GPU else "cv2-Telea(CPU)"
    status    = f"✅ Done in {elapsed}s | {seg_tag} + {fill_tag}"

    return subject_img, filled_bg, composite, status

# ─────────────────────────────────────────────────────────────────────────────
# Admin helpers
# ─────────────────────────────────────────────────────────────────────────────
_admin_auth = {"logged_in": False}

def admin_login(password: str):
    if password == ADMIN_PASSWORD:
        _admin_auth["logged_in"] = True
        return gr.update(visible=False), gr.update(visible=True), "✅ Logged in successfully."
    return gr.update(visible=True), gr.update(visible=False), "❌ Wrong password."

def _require_auth():
    if not _admin_auth["logged_in"]:
        raise PermissionError("Not logged in")

def _customers_markdown() -> str:
    _reset_if_new_day()
    if not API_KEYS:
        return "_No customers yet._"
    rows = [
        "| # | Email | Plan | Key (preview) | Calls Today | Limit | Created | Active |",
        "|---|-------|------|---------------|-------------|-------|---------|--------|",
    ]
    for i, (key, info) in enumerate(API_KEYS.items(), 1):
        plan    = info.get("plan", "free")
        limit   = info.get("custom_limit") or PLAN_LIMITS.get(plan, 10)
        calls   = _used_today(key)
        preview = key[:22] + "…"
        created = info.get("created", "—")
        active  = "✅" if info.get("active", True) else "❌"
        email   = info.get("email", "—")
        rows.append(
            f"| {i} | {email} | **{plan}** | `{preview}` | {calls} | {limit} | {created} | {active} |"
        )
    return "\n".join(rows)

def admin_add_customer(email: str, plan: str, custom_limit_str: str):
    try:
        _require_auth()
    except PermissionError:
        return "❌ Not logged in.", _customers_markdown()
    email = email.strip()
    if not email:
        return "❌ Email is required.", _customers_markdown()
    new_key = "freemg-" + uuid.uuid4().hex[:16]
    entry   = {"email": email, "plan": plan, "created": str(date.today()), "active": True}
    if custom_limit_str.strip():
        try:
            entry["custom_limit"] = int(custom_limit_str.strip())
        except ValueError:
            return "❌ Custom limit must be a number.", _customers_markdown()
    API_KEYS[new_key] = entry
    _persist_keys()
    limit_display = entry.get("custom_limit") or PLAN_LIMITS.get(plan, 10)
    return (
        f"✅ **Key created!**\n\n```\n{new_key}\n```\n\n"
        f"**Email:** {email}  |  **Plan:** {plan}  |  **Daily limit:** {limit_display}",
        _customers_markdown()
    )

def admin_upgrade_plan(key: str, new_plan: str):
    try:
        _require_auth()
    except PermissionError:
        return "❌ Not logged in.", _customers_markdown()
    key = key.strip()
    if key not in API_KEYS:
        return "❌ Key not found.", _customers_markdown()
    old_plan = API_KEYS[key].get("plan", "free")
    API_KEYS[key]["plan"] = new_plan
    API_KEYS[key].pop("custom_limit", None)
    _persist_keys()
    return f"✅ `{key[:22]}…` upgraded **{old_plan}** → **{new_plan}**", _customers_markdown()

def admin_delete_key(key: str):
    try:
        _require_auth()
    except PermissionError:
        return "❌ Not logged in.", _customers_markdown()
    key = key.strip()
    if key not in API_KEYS:
        return "❌ Key not found.", _customers_markdown()
    email = API_KEYS[key].get("email", "—")
    del API_KEYS[key]
    _persist_keys()
    return f"✅ Deleted key for **{email}**", _customers_markdown()

def admin_toggle_active(key: str):
    try:
        _require_auth()
    except PermissionError:
        return "❌ Not logged in.", _customers_markdown()
    key = key.strip()
    if key not in API_KEYS:
        return "❌ Key not found.", _customers_markdown()
    current = API_KEYS[key].get("active", True)
    API_KEYS[key]["active"] = not current
    _persist_keys()
    return f"✅ Key `{key[:22]}…` {'activated' if not current else 'deactivated'}", _customers_markdown()

def admin_set_custom_limit(key: str, new_limit_str: str):
    try:
        _require_auth()
    except PermissionError:
        return "❌ Not logged in.", _customers_markdown()
    key = key.strip()
    if key not in API_KEYS:
        return "❌ Key not found.", _customers_markdown()
    try:
        new_limit = int(new_limit_str.strip())
    except ValueError:
        return "❌ Limit must be an integer.", _customers_markdown()
    API_KEYS[key]["custom_limit"] = new_limit
    _persist_keys()
    return f"✅ Limit for `{key[:22]}…` set to **{new_limit}/day**", _customers_markdown()

def admin_stats() -> str:
    _reset_if_new_day()
    total      = len(API_KEYS)
    by_plan    = {}
    revenue    = 0
    for info in API_KEYS.values():
        p = info.get("plan", "free")
        by_plan[p] = by_plan.get(p, 0) + 1
        revenue   += PLAN_PRICES.get(p, 0)
    total_calls = sum(_usage_today.values())
    plan_rows   = "\n".join(
        f"  - **{p}**: {c} customer{'s' if c!=1 else ''} × ${PLAN_PRICES.get(p,0)}/mo = **${PLAN_PRICES.get(p,0)*c}/mo**"
        for p, c in sorted(by_plan.items())
    )
    mode_tag = "🟢 GPU (SAM2.1 + SDXL)" if HAS_GPU else "🟡 CPU (GrabCut + cv2)"
    return f"""### 📊 Live Stats — {date.today()}

| Metric | Value |
|--------|-------|
| Total Customers | **{total}** |
| API Calls Today | **{total_calls}** |
| Estimated MRR | **${revenue}/mo** |
| Space Mode | {mode_tag} |

**Revenue Breakdown:**
{plan_rows if plan_rows else "  _No customers yet._"}
"""

def admin_export_json() -> str:
    return json.dumps({
        "exported_at": str(datetime.now()),
        "repo":        HF_REPO_ID,
        "total_keys":  len(API_KEYS),
        "customers":   API_KEYS,
        "usage_today": _usage_today,
    }, indent=2)

# ─────────────────────────────────────────────────────────────────────────────
# CSS
# ─────────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700;800&family=Space+Mono:wght@400;700&display=swap');
:root {
    --bg:#080808; --surf:#111; --surf2:#181818; --border:#252525;
    --acc:#FF6B35; --acc2:#FFD700; --green:#22c55e; --red:#ef4444;
    --blue:#818cf8; --text:#efefef; --muted:#777; --r:10px;
    --font:'Syne',sans-serif; --mono:'Space Mono',monospace;
}
body,.gradio-container{background:var(--bg)!important;color:var(--text)!important;font-family:var(--font)!important;}
.tabs>.tab-nav{background:var(--surf);border-bottom:1px solid var(--border);border-radius:var(--r) var(--r) 0 0;padding:0 8px;}
.tabs>.tab-nav button{color:var(--muted)!important;font-family:var(--font)!important;font-weight:700;padding:12px 20px;border-bottom:2px solid transparent;transition:all .2s;border-radius:0!important;background:transparent!important;}
.tabs>.tab-nav button.selected{color:var(--acc)!important;border-bottom-color:var(--acc)!important;}
button.primary,.gr-button-primary{background:linear-gradient(135deg,var(--acc),#e85d2a)!important;color:#fff!important;border:none!important;font-family:var(--font)!important;font-weight:700!important;border-radius:8px!important;transition:opacity .2s,transform .1s!important;}
button.primary:hover{opacity:.88!important;transform:translateY(-1px)!important;}
button.secondary{background:var(--surf2)!important;color:var(--text)!important;border:1px solid var(--border)!important;font-family:var(--font)!important;font-weight:600!important;border-radius:8px!important;}
button.stop,.gr-button-stop{background:var(--red)!important;color:#fff!important;border:none!important;font-family:var(--font)!important;font-weight:700!important;border-radius:8px!important;}
input,textarea,select{background:var(--surf2)!important;border:1px solid var(--border)!important;color:var(--text)!important;border-radius:8px!important;font-family:var(--font)!important;transition:border-color .2s;}
input:focus,textarea:focus{border-color:var(--acc)!important;outline:none!important;box-shadow:0 0 0 2px rgba(255,107,53,.15)!important;}
.gr-box,.gr-panel{background:var(--surf)!important;border:1px solid var(--border)!important;border-radius:var(--r)!important;}
label{color:var(--muted)!important;font-weight:700!important;font-size:.8rem!important;text-transform:uppercase;letter-spacing:.5px;}
.gr-markdown{color:var(--text)!important;}
.gr-markdown code{background:var(--surf2);border:1px solid var(--border);border-radius:5px;padding:1px 6px;font-family:var(--mono);font-size:.85em;color:var(--acc2);}
.gr-markdown pre{background:var(--surf2);border:1px solid var(--border);border-radius:8px;padding:16px;overflow-x:auto;}
.gr-markdown table{width:100%;border-collapse:collapse;}
.gr-markdown th{background:var(--surf2);color:var(--muted);font-size:.75rem;text-transform:uppercase;padding:10px 14px;border-bottom:1px solid var(--border);text-align:left;}
.gr-markdown td{padding:10px 14px;border-bottom:1px solid #1a1a1a;color:var(--text);}
.gr-markdown tr:last-child td{border-bottom:none;}
.mg-header{background:linear-gradient(135deg,#130a05 0%,#0d0d0d 60%,#09120d 100%);border:1px solid var(--border);border-top:3px solid var(--acc);border-radius:var(--r);padding:28px 36px 24px;margin-bottom:20px;position:relative;overflow:hidden;}
.mg-header::after{content:'✨';position:absolute;right:32px;top:50%;transform:translateY(-50%);font-size:5rem;opacity:.06;pointer-events:none;}
.mg-title{font-size:1.9rem;font-weight:800;margin:0 0 4px;background:linear-gradient(90deg,var(--acc),var(--acc2));-webkit-background-clip:text;-webkit-text-fill-color:transparent;}
.mg-sub{color:var(--muted);font-size:.88rem;margin:0;}
.mg-sub a{color:var(--acc);text-decoration:none;}
.plan-table{width:100%;border-collapse:collapse;border-radius:10px;overflow:hidden;border:1px solid var(--border);}
.plan-table th{background:var(--surf2);color:var(--muted);font-size:.75rem;text-transform:uppercase;letter-spacing:.5px;padding:12px 16px;text-align:left;border-bottom:1px solid var(--border);}
.plan-table td{padding:12px 16px;border-bottom:1px solid #1a1a1a;font-family:var(--mono);font-size:.88rem;}
.plan-table tr:last-child td{border-bottom:none;}
.badge{display:inline-block;padding:2px 10px;border-radius:20px;font-size:.72rem;font-weight:700;text-transform:uppercase;letter-spacing:.4px;font-family:var(--font);}
.badge-free{background:#252525;color:#aaa;}
.badge-starter{background:#0d2b1a;color:#4ade80;}
.badge-pro{background:#14103a;color:#818cf8;}
.price-free{color:var(--green);font-weight:700;}
.price-starter{color:var(--blue);font-weight:700;}
.price-pro{color:var(--acc2);font-weight:700;}
.mode-banner{background:var(--surf2);border:1px solid var(--border);border-radius:8px;padding:10px 16px;margin-bottom:12px;font-family:var(--mono);font-size:.82rem;color:var(--muted);}
@keyframes dotPop{from{transform:translate(-50%,-50%) scale(0);opacity:0;}to{transform:translate(-50%,-50%) scale(1);opacity:1;}}
@keyframes dotFade{from{opacity:1;}to{opacity:0;}}
"""

# ─────────────────────────────────────────────────────────────────────────────
# JS — click-point collector
# ─────────────────────────────────────────────────────────────────────────────
CLICK_JS = """
<script>
(function(){
  var pts=[];
  function getField(){ return document.querySelector('#mg_points_field textarea'); }
  function syncField(){
    var f=getField(); if(!f)return;
    f.value=JSON.stringify(pts);
    f.dispatchEvent(new Event('input',{bubbles:true}));
    f.dispatchEvent(new Event('change',{bubbles:true}));
  }
  function showDot(cx,cy,color){
    var d=document.createElement('div');
    d.style.cssText='position:fixed;width:14px;height:14px;border-radius:50%;background:'+color+
      ';border:2px solid #fff;pointer-events:none;z-index:99999;left:'+cx+'px;top:'+cy+'px;'+
      'transform:translate(-50%,-50%);animation:dotPop .25s ease forwards;';
    document.body.appendChild(d);
    setTimeout(function(){d.style.animation='dotFade .4s ease forwards';setTimeout(function(){d.remove();},500);},2500);
  }
  function attach(){
    document.querySelectorAll('.mg-canvas').forEach(function(c){
      if(c._mga)return; c._mga=true;
      c.addEventListener('click',function(e){
        var img=e.target.closest('img'); if(!img)return;
        var r=img.getBoundingClientRect();
        var sx=(img.naturalWidth||img.width)/r.width;
        var sy=(img.naturalHeight||img.height)/r.height;
        var ix=(e.clientX-r.left)*sx, iy=(e.clientY-r.top)*sy;
        var label=e.shiftKey?0:1;
        pts.push({x:Math.round(ix),y:Math.round(iy),label:label});
        syncField();
        showDot(e.clientX,e.clientY,label===1?'#22c55e':'#ef4444');
      });
    });
  }
  window.mgClearPoints=function(){pts=[];syncField();};
  var poll=setInterval(function(){
    if(document.querySelector('.mg-canvas')){
      attach();
      new MutationObserver(attach).observe(document.body,{childList:true,subtree:true});
      clearInterval(poll);
    }
  },1000);
})();
</script>
"""

PLANS_HTML = """
<div style="margin-top:24px;">
<table class="plan-table">
<thead><tr><th>Plan</th><th>Daily Limit</th><th>Models</th><th>Price</th></tr></thead>
<tbody>
<tr><td><span class="badge badge-free">Free</span></td><td>10 / day</td><td>fast (cv2)</td><td class="price-free">$0</td></tr>
<tr><td><span class="badge badge-starter">Starter</span></td><td>100 / day</td><td>fast + quality</td><td class="price-starter">$9 / mo</td></tr>
<tr><td><span class="badge badge-pro">Pro</span></td><td>500 / day</td><td>All 3 models</td><td class="price-pro">$29 / mo</td></tr>
</tbody></table>
<p style="color:#555;font-size:.78rem;margin:10px 0 0;font-family:'Space Mono',monospace;">
📧 <a href="mailto:support@freebg.site" style="color:#FF6B35;text-decoration:none;">support@freebg.site</a>
&nbsp;·&nbsp; 🔑 Master key = unlimited calls for your own website
</p></div>
"""

# ─────────────────────────────────────────────────────────────────────────────
# Build Gradio UI
# ─────────────────────────────────────────────────────────────────────────────
def build_ui() -> gr.Blocks:
    mode_banner_html = (
        '<div class="mode-banner">⚡ Running on <strong style="color:#22c55e;">GPU</strong> — SAM2.1 + SDXL inpainting active</div>'
        if HAS_GPU else
        '<div class="mode-banner">🟡 Running on <strong style="color:#FFD700;">CPU</strong> — GrabCut + cv2 segmentation (fast, ~3s) | GPU grant pending</div>'
    )

    with gr.Blocks(css=CSS, title="FreeMG — Magic Grab API", theme=gr.themes.Base()) as demo:

        gr.HTML(CLICK_JS)

        gr.HTML("""
        <div class="mg-header">
            <p class="mg-title">✨ FreeMG — Magic Grab API</p>
            <p class="mg-sub">SAM2.1 · GrabCut · SDXL Inpainting &nbsp;·&nbsp; API Key Auth &nbsp;·&nbsp;
               <a href="https://freebg.site">freebg.site</a></p>
        </div>
        """)

        with gr.Tabs():

            # ══ Tab 1: Magic Grab ═════════════════════════════════════════════
            with gr.Tab("✨ Magic Grab"):
                gr.HTML(mode_banner_html)
                gr.Markdown(
                    "> **Click** on your subject → green dot. **Shift+click** = background hint.  \n"
                    "> Then hit **Grab Now!**"
                )

                with gr.Row(equal_height=False):
                    with gr.Column(scale=1, min_width=340):
                        img_in = gr.Image(
                            label="📷 Upload Image — click to select subject",
                            type="pil", height=380,
                            elem_classes=["mg-canvas"],
                        )
                        pts_field = gr.Textbox(
                            value="[]", label="Click points (auto-filled by JS)",
                            elem_id="mg_points_field",
                            info='Click on image above, or type JSON: [{"x":320,"y":240,"label":1}]',
                            lines=2,
                        )
                        mode_radio = gr.Radio(
                            choices=["fast", "quality", "best"], value="fast",
                            label="⚙️ Mode",
                            info="fast=cv2(CPU,instant) | quality=SDXL 20steps(GPU) | best=SDXL 35steps(GPU)",
                        )
                        fill_prompt_in = gr.Textbox(
                            label="🎨 Background Fill Prompt (quality/best + GPU only)",
                            placeholder="e.g. sunny outdoor park with grass",
                            lines=2,
                        )
                        output_type_radio = gr.Radio(
                            choices=["subject_only", "filled_bg", "composite"],
                            value="composite", label="📤 Primary Output",
                        )
                        api_key_in = gr.Textbox(
                            label="🔑 API Key",
                            placeholder="freemg-xxxx  (blank = UI demo mode)",
                            type="password",
                        )
                        with gr.Row():
                            clear_btn = gr.Button("🗑️ Clear Points", variant="secondary")
                            grab_btn  = gr.Button("✨ Grab Now!", variant="primary", scale=2)

                    with gr.Column(scale=1, min_width=340):
                        out_subject   = gr.Image(label="✂️ Subject (Transparent PNG)", type="pil")
                        out_bg_filled = gr.Image(label="🖼️ Filled Background",         type="pil")
                        out_composite = gr.Image(label="🎯 Final Composite",           type="pil")
                        status_out    = gr.Textbox(label="⚡ Status", interactive=False)

                grab_btn.click(
                    fn=magic_grab_pipeline,
                    inputs=[img_in, pts_field, mode_radio, fill_prompt_in,
                            output_type_radio, api_key_in],
                    outputs=[out_subject, out_bg_filled, out_composite, status_out],
                )
                clear_btn.click(fn=None, js="()=>{if(window.mgClearPoints)window.mgClearPoints();}")
                gr.HTML(PLANS_HTML)

            # ══ Tab 2: My Usage ═══════════════════════════════════════════════
            with gr.Tab("📊 My Usage"):
                gr.Markdown("### Check Your API Key Usage")
                usage_key = gr.Textbox(label="Your API Key", type="password", placeholder="freemg-xxxx")
                check_btn = gr.Button("Check Usage", variant="primary")
                usage_md  = gr.Markdown()

                def _check_usage(key):
                    _reset_if_new_day()
                    _, reason, info = validate_key(key)
                    if reason == "no_key":
                        return "⚠️ Please enter your API key."
                    if reason == "invalid_key":
                        return "❌ **Invalid API key.** Contact support@freebg.site"
                    plan  = info.get("plan", "free")
                    limit = info.get("custom_limit") or PLAN_LIMITS.get(plan, 10)
                    used  = _used_today(key.strip())
                    rem   = max(0, limit - used)
                    pct   = min(100, round((used / limit) * 100)) if limit else 0
                    bar   = "█" * (pct // 5) + "░" * (20 - pct // 5)
                    color = "🟢" if pct < 70 else ("🟡" if pct < 90 else "🔴")
                    return (
                        f"**Email:** {info.get('email','—')}  \n"
                        f"**Plan:** `{plan}`  \n"
                        f"**Daily Limit:** {limit}  \n"
                        f"**Used Today:** {used}  \n"
                        f"**Remaining:** {rem}  \n\n"
                        f"{color} `{bar}` {pct}%  \n\n"
                        f"**Created:** {info.get('created','—')}  \n"
                        f"**Status:** {'✅ Active' if info.get('active',True) else '❌ Inactive'}"
                    )

                check_btn.click(fn=_check_usage, inputs=[usage_key], outputs=[usage_md])

            # ══ Tab 3: API Docs ═══════════════════════════════════════════════
            with gr.Tab("📄 API Docs"):
                gr.Markdown("""
## FreeMG REST API

**Base URL:** `https://freebg-magic-grab.hf.space`

---

### `POST /api/magic-grab`

| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `api_key` | string | ✅ | Your FreeMG API key |
| `image` | file | ✅ | JPEG or PNG, max 10 MB |
| `mode` | string | — | `fast` / `quality` / `best` |
| `fill_prompt` | string | — | Background fill text prompt |
| `output_type` | string | — | `subject_only` / `filled_bg` / `composite` |
| `click_x` | string | — | Comma-separated X coords e.g. `"320,400"` |
| `click_y` | string | — | Comma-separated Y coords e.g. `"240,210"` |

**Success Response:**
```json
{
  "success": true,
  "subject": "<base64 PNG transparent>",
  "filled_bg": "<base64 PNG>",
  "composite": "<base64 PNG>",
  "elapsed": 1.82,
  "model": "GrabCut(CPU) + cv2-Telea"
}
```

**Error Response:**
```json
{ "success": false, "error": "quota_exceeded | invalid_key | processing_error" }
```

---

### Python Example
```python
import requests, base64
from PIL import Image
from io import BytesIO

resp = requests.post(
    "https://freebg-magic-grab.hf.space/api/magic-grab",
    data={
        "api_key":     "freemg-your-key",
        "mode":        "fast",
        "click_x":     "320",
        "click_y":     "240",
    },
    files={"image": open("photo.jpg", "rb")},
)
data = resp.json()
if data["success"]:
    img = Image.open(BytesIO(base64.b64decode(data["composite"])))
    img.save("result.png")
```

---

### JavaScript
```javascript
const form = new FormData();
form.append("api_key", "freemg-your-key");
form.append("mode",    "fast");
form.append("click_x", "320");
form.append("click_y", "240");
form.append("image",   fileInput.files[0]);

const resp = await fetch(
    "https://freebg-magic-grab.hf.space/api/magic-grab",
    { method: "POST", body: form }
);
const data = await resp.json();
if (data.success) {
    imgEl.src = "data:image/png;base64," + data.composite;
}
```
                """)

            # ══ Tab 4: Admin ══════════════════════════════════════════════════
            with gr.Tab("🔐 Admin"):

                with gr.Group(visible=True) as login_group:
                    gr.Markdown("### 🔒 Admin Login")
                    pwd_in    = gr.Textbox(label="Password", type="password")
                    login_btn = gr.Button("Login", variant="primary")
                    login_msg = gr.Markdown()

                with gr.Group(visible=False) as dashboard_group:
                    gr.Markdown("## 👑 Admin Dashboard")

                    with gr.Tabs():

                        with gr.Tab("👥 Customers"):
                            refresh_cust = gr.Button("🔄 Refresh", variant="secondary")
                            cust_md = gr.Markdown(value=_customers_markdown())
                            refresh_cust.click(fn=_customers_markdown, outputs=[cust_md])

                        with gr.Tab("➕ Add Customer"):
                            gr.Markdown("### Create New API Key")
                            with gr.Row():
                                add_email = gr.Textbox(label="Email", placeholder="user@example.com")
                                add_plan  = gr.Dropdown(choices=["free","starter","pro"],
                                                        value="free", label="Plan")
                            add_limit = gr.Textbox(label="Custom Daily Limit (optional)",
                                                   placeholder="blank = plan default")
                            add_btn  = gr.Button("✅ Create Key + Save to HF Secret", variant="primary")
                            add_msg  = gr.Markdown()
                            add_tbl  = gr.Markdown()
                            add_btn.click(fn=admin_add_customer,
                                          inputs=[add_email, add_plan, add_limit],
                                          outputs=[add_msg, add_tbl])

                        with gr.Tab("⚙️ Manage Keys"):
                            gr.Markdown("#### ⬆️ Upgrade Plan")
                            with gr.Row():
                                upg_key  = gr.Textbox(label="Full API Key")
                                upg_plan = gr.Dropdown(choices=["free","starter","pro"], label="New Plan")
                            upg_btn = gr.Button("Apply Plan Change", variant="primary")
                            upg_msg = gr.Markdown()
                            upg_tbl = gr.Markdown()
                            upg_btn.click(fn=admin_upgrade_plan,
                                          inputs=[upg_key, upg_plan],
                                          outputs=[upg_msg, upg_tbl])

                            gr.Markdown("---")
                            gr.Markdown("#### 🔢 Set Custom Limit")
                            with gr.Row():
                                lim_key = gr.Textbox(label="Full API Key")
                                lim_val = gr.Textbox(label="New Limit", placeholder="e.g. 250")
                            lim_btn = gr.Button("Set Limit", variant="primary")
                            lim_msg = gr.Markdown()
                            lim_tbl = gr.Markdown()
                            lim_btn.click(fn=admin_set_custom_limit,
                                          inputs=[lim_key, lim_val],
                                          outputs=[lim_msg, lim_tbl])

                            gr.Markdown("---")
                            gr.Markdown("#### ⏸️ Toggle Active")
                            tog_key = gr.Textbox(label="Full API Key")
                            tog_btn = gr.Button("Toggle Active/Inactive", variant="secondary")
                            tog_msg = gr.Markdown()
                            tog_tbl = gr.Markdown()
                            tog_btn.click(fn=admin_toggle_active,
                                          inputs=[tog_key],
                                          outputs=[tog_msg, tog_tbl])

                            gr.Markdown("---")
                            gr.Markdown("#### 🗑️ Delete Key")
                            del_key = gr.Textbox(label="Full API Key to Delete")
                            del_btn = gr.Button("Delete Key Permanently", variant="stop")
                            del_msg = gr.Markdown()
                            del_tbl = gr.Markdown()
                            del_btn.click(fn=admin_delete_key,
                                          inputs=[del_key],
                                          outputs=[del_msg, del_tbl])

                        with gr.Tab("📈 Stats"):
                            stats_md    = gr.Markdown(value=admin_stats())
                            refresh_st  = gr.Button("🔄 Refresh", variant="secondary")
                            refresh_st.click(fn=admin_stats, outputs=[stats_md])

                        with gr.Tab("💾 Export JSON"):
                            gr.Markdown("Full backup of customer keys + today's usage.")
                            exp_btn = gr.Button("📥 Export JSON", variant="primary")
                            exp_out = gr.Code(language="json", label="Data")
                            exp_btn.click(fn=admin_export_json, outputs=[exp_out])

                        with gr.Tab("🔑 HF Secrets"):
                            gr.Markdown("""
### Required Secrets

| Secret | Value |
|--------|-------|
| `MASTER_API_KEY` | Your internal unlimited key |
| `API_KEYS_JSON` | `{}` — auto-managed by admin |
| `HF_TOKEN` | Write-access HF token |
| `HF_REPO_ID` | `freebg/magic-grab` |
| `ADMIN_PASSWORD` | Admin panel password |

> Admin panel auto-saves `API_KEYS_JSON` to HF secret on every change.
                            """)
                            sync_btn = gr.Button("Force Sync Keys to HF Secret", variant="secondary")
                            sync_msg = gr.Markdown()
                            sync_btn.click(
                                fn=lambda: "✅ Synced." if _persist_keys() else "⚠️ Failed — check HF_TOKEN.",
                                outputs=[sync_msg]
                            )

                login_btn.click(
                    fn=admin_login,
                    inputs=[pwd_in],
                    outputs=[login_group, dashboard_group, login_msg],
                )

    return demo

# ─────────────────────────────────────────────────────────────────────────────
# REST API endpoint
# ─────────────────────────────────────────────────────────────────────────────
def mount_rest_api(demo: gr.Blocks):
    @demo.app.post("/api/magic-grab")
    async def _api(
        api_key:     str        = Form(...),
        image:       UploadFile = File(...),
        mode:        str        = Form("fast"),
        fill_prompt: str        = Form(""),
        output_type: str        = Form("composite"),
        click_x:     str        = Form(""),
        click_y:     str        = Form(""),
    ):
        try:
            valid, reason, _ = validate_key(api_key)
            if not valid:
                return JSONResponse({"success": False, "error": reason}, status_code=403)

            raw = await image.read()
            pil = Image.open(BytesIO(raw)).convert("RGB")

            if click_x.strip() and click_y.strip():
                xs  = [v.strip() for v in click_x.split(",") if v.strip()]
                ys  = [v.strip() for v in click_y.split(",") if v.strip()]
                pts = [{"x": int(x), "y": int(y), "label": 1} for x, y in zip(xs, ys)]
            else:
                pts = [{"x": pil.width // 2, "y": pil.height // 2, "label": 1}]

            subject, filled, composite, status = magic_grab_pipeline(
                pil, json.dumps(pts), mode, fill_prompt, output_type, api_key
            )

            result = {"success": True, "status": status}
            if subject:   result["subject"]   = _pil_to_b64(subject)
            if filled:    result["filled_bg"]  = _pil_to_b64(filled)
            if composite: result["composite"]  = _pil_to_b64(composite)
            return JSONResponse(result)

        except Exception as exc:
            print(f"[API] {exc}\n{traceback.format_exc()}")
            return JSONResponse({"success": False, "error": str(exc)}, status_code=500)

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo = build_ui()
    demo.queue(max_size=10)
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
else:
    demo = build_ui()
    demo.queue(max_size=10)
    try:
        mount_rest_api(demo)
        print("[REST] /api/magic-grab endpoint mounted ✓")
    except Exception as e:
        print(f"[REST] Mount failed: {e}")