import gradio as gr
import numpy as np
from PIL import Image
import io
import json
import os
import time
import secrets
import datetime
from fastapi import Request
from fastapi.responses import JSONResponse, Response

# ── ZeroGPU / SD1.5 auto-detect ──────────────────────────────────────────────
HAS_ZERO_GPU = False
HAS_CUDA = False

try:
    import spaces
    # Only enable if actually running on ZeroGPU hardware
    import os
    if os.environ.get("SPACE_HARDWARE") in ("gpu-t4-s", "gpu-t4-m", "gpu-a10g-s", "gpu-a10g-l", "zero-gpu"):
        HAS_ZERO_GPU = True
        print("✅ ZeroGPU detected — SD1.5 'best' mode will be available")
    else:
        print("ℹ️  spaces imported but no ZeroGPU hardware — LaMa-only mode")
except ImportError:
    print("ℹ️  No ZeroGPU — LaMa-only mode (fast + quality)")

try:
    import torch
    HAS_CUDA = torch.cuda.is_available()
except:
    pass

# ── HF Secret auto-update (like FreeBG) ──────────────────────────────────────
HF_TOKEN   = os.environ.get("HF_TOKEN", "")
HF_REPO_ID = os.environ.get("HF_REPO_ID", "")   # e.g. "freebg/magic-eraser"

def save_keys_to_hf(keys: dict) -> bool:
    """Push updated API_KEYS_JSON straight to HF Space Secret — Space restarts automatically."""
    if not HF_TOKEN or not HF_REPO_ID:
        return False
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        api.add_space_secret(
            repo_id=HF_REPO_ID,
            key="API_KEYS_JSON",
            value=json.dumps(keys),
        )
        return True
    except Exception as e:
        print(f"⚠️  HF Secret update failed: {e}")
        return False

# ── LaMa model ────────────────────────────────────────────────────────────────
lama_model = None

def load_lama():
    global lama_model
    if lama_model is not None:
        return lama_model
    try:
        from simple_lama_inpainting import SimpleLama
        lama_model = SimpleLama()
        print("✅ LaMa loaded")
        return lama_model
    except Exception as e:
        print(f"❌ LaMa load error: {e}")
        return None

# ── SD1.5 pipeline (ZeroGPU / CUDA only) ─────────────────────────────────────
sd_pipe = None

def load_sd():
    return None

# ── Key / Plan config ─────────────────────────────────────────────────────────
MASTER_KEY     = os.environ.get("MASTER_API_KEY", "freebg-master-change-this")
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "admin123")

PLAN_LIMITS = {
    "free":    {"daily": 10,    "models": ["fast"]},
    "starter": {"daily": 100,   "models": ["fast", "quality"]},
    "pro":     {"daily": 500,   "models": ["fast", "quality", "best"]},
    "master":  {"daily": 99999, "models": ["fast", "quality", "best"]},
}

# In-memory keys cache (reloaded from env each call so secrets take effect after restart)
usage_cache: dict = {}

def get_today() -> str:
    return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")

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

def check_and_count(api_key: str, model: str) -> tuple[bool, str]:
    if api_key == MASTER_KEY:
        return True, ""
    keys = load_keys()
    if api_key not in keys:
        return False, "Invalid API key. Get one at freebg.site"
    info  = keys[api_key]
    plan  = info.get("plan", "free")
    lims  = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
    if model not in lims["models"]:
        return False, f"Model '{model}' not available on {plan} plan. Upgrade at freebg.site"
    today     = get_today()
    cache_key = f"{api_key}:{today}"
    used      = usage_cache.get(cache_key, 0)
    if used >= lims["daily"]:
        return False, f"Daily limit reached ({lims['daily']}/day on {plan} plan)"
    usage_cache[cache_key] = used + 1
    return True, ""

def get_usage_info(api_key: str) -> dict:
    """Always returns a dict with 'models' key — no KeyError possible."""
    plan_info = PLAN_LIMITS.get("master")
    if api_key == MASTER_KEY:
        today = get_today()
        used  = usage_cache.get(f"{api_key}:{today}", 0)
        return {
            "plan":        "master",
            "daily_limit": plan_info["daily"],
            "used_today":  used,
            "remaining":   plan_info["daily"] - used,
            "models":      plan_info["models"],
            "owner":       "freebg.site",
        }
    keys = load_keys()
    if api_key not in keys:
        return {"error": "Invalid API key"}
    info  = keys[api_key]
    plan  = info.get("plan", "free")
    lims  = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
    today = get_today()
    used  = usage_cache.get(f"{api_key}:{today}", 0)
    return {
        "plan":        plan,
        "daily_limit": lims["daily"],
        "used_today":  used,
        "remaining":   max(0, lims["daily"] - used),
        "models":      lims["models"],
        "owner":       info.get("owner", ""),
    }

# ── Inpainting ────────────────────────────────────────────────────────────────
def _prep_mask(mask: Image.Image) -> Image.Image:
    arr = np.array(mask.convert("L"))
    return Image.fromarray(((arr > 127).astype(np.uint8) * 255), "L")

def run_lama(image: Image.Image, mask: Image.Image) -> Image.Image:
    lama = load_lama()
    if lama is None:
        raise RuntimeError("LaMa model failed to load")
    return lama(image.convert("RGB"), _prep_mask(mask))

def run_lama_quality(image: Image.Image, mask: Image.Image) -> Image.Image:
    from PIL import ImageFilter
    return run_lama(image, mask).filter(ImageFilter.UnsharpMask(radius=1.5, percent=120, threshold=3))

def run_best(image: Image.Image, mask: Image.Image) -> Image.Image:
    return run_lama_quality(image, mask)

if HAS_ZERO_GPU:
    import spaces as _spaces
    run_lama         = _spaces.GPU(run_lama)
    run_lama_quality = _spaces.GPU(run_lama_quality)
    run_best         = _spaces.GPU(run_best)

# ── Gradio helpers ────────────────────────────────────────────────────────────
def erase_ui(composite, model: str, api_key: str):
    if composite is None:
        return None, "❌ Please upload and paint an image first"
    background = composite.get("background")
    layers     = composite.get("layers", [])
    if background is None:
        return None, "❌ No image provided"
    image = background.convert("RGB")
    if not layers:
        return None, "❌ Please paint over the area to erase"
    r, g, b, a = layers[0].convert("RGBA").split()
    mask = a
    if np.array(mask).max() == 0:
        return None, "❌ No brush strokes found — paint over the area to erase"
    key     = (api_key or "").strip() or MASTER_KEY
    allowed, err = check_and_count(key, model.lower())
    if not allowed:
        return None, f"❌ {err}"
    try:
        t0     = time.time()
        fn     = {"fast": run_lama, "quality": run_lama_quality, "best": run_best}.get(model, run_lama)
        result = fn(image, mask)
        return result, f"✅ Done in {time.time()-t0:.1f}s  |  Model: {model}"
    except Exception as e:
        return None, f"❌ Error: {e}"

def check_usage_ui(api_key: str) -> str:
    key = (api_key or "").strip()
    if not key:
        return "❌ Enter your API key above"
    info = get_usage_info(key)
    if "error" in info:
        return f"❌ {info['error']}"
    return (
        f"**Plan:** {info['plan'].upper()}\n\n"
        f"**Calls today:** {info['used_today']} / {info['daily_limit']}\n\n"
        f"**Remaining:** {info['remaining']}\n\n"
        f"**Models:** {', '.join(info['models'])}\n\n"
        f"**Owner:** {info['owner']}"
    )

# ── Admin helpers ─────────────────────────────────────────────────────────────
def _auth(pw: str) -> bool:
    return pw == ADMIN_PASSWORD

def admin_add_customer(pw, email, plan, custom_limit):
    if not _auth(pw):
        return "❌ Wrong password", "", ""
    if plan not in PLAN_LIMITS:
        return f"❌ Invalid plan", "", ""
    prefix = {"free":"me-fr","starter":"me-st","pro":"me-pr","master":"me-ma"}[plan]
    key    = f"freebg-{prefix}-{secrets.token_hex(4)}-{secrets.token_hex(4)}"
    keys   = load_keys()
    entry  = {"owner": email, "plan": plan, "created": get_today(), "active": True}
    if custom_limit:
        try:
            entry["custom_daily"] = int(custom_limit)
        except:
            pass
    keys[key] = entry
    saved = save_keys_to_hf(keys)
    hf_note = "✅ Auto-saved to HF Secret. Space restarts in ~30s." if saved else "⚠️ HF auto-save failed — paste JSON manually."
    return f"✅ Key created for {email} ({plan})\n\nKey: **{key}**\n\n{hf_note}", key, json.dumps(keys, indent=2)

def admin_get_customers(pw):
    if not _auth(pw):
        return "❌ Wrong password"
    keys  = load_keys()
    today = get_today()
    rows  = []
    for k, v in keys.items():
        used  = usage_cache.get(f"{k}:{today}", 0)
        plan  = v.get("plan","free")
        limit = v.get("custom_daily", PLAN_LIMITS.get(plan,{}).get("daily","?"))
        rows.append(f"| `{k}` | {v.get('owner','')} | {plan} | {used} | {limit} | {v.get('created','')} |")
    if not rows:
        return "No customers yet."
    header = "| API Key | Owner | Plan | Calls Today | Daily Limit | Created |\n|---|---|---|---|---|---|"
    return header + "\n" + "\n".join(rows)

def admin_upgrade(pw, key, new_plan):
    if not _auth(pw):
        return "❌ Wrong password", ""
    keys = load_keys()
    if key not in keys:
        return "❌ Key not found", ""
    old = keys[key].get("plan","?")
    keys[key]["plan"] = new_plan
    saved = save_keys_to_hf(keys)
    note  = "✅ Auto-saved to HF Secret." if saved else "⚠️ Paste JSON manually."
    return f"✅ {key}\n{old} → {new_plan}\n{note}", json.dumps(keys, indent=2)

def admin_delete(pw, key):
    if not _auth(pw):
        return "❌ Wrong password", ""
    keys = load_keys()
    if key not in keys:
        return "❌ Key not found", ""
    del keys[key]
    saved = save_keys_to_hf(keys)
    note  = "✅ Auto-saved to HF Secret." if saved else "⚠️ Paste JSON manually."
    return f"✅ Deleted: {key}\n{note}", json.dumps(keys, indent=2)

def admin_stats(pw):
    if not _auth(pw):
        return "❌ Wrong password"
    keys  = load_keys()
    today = get_today()
    plans = {"free":0,"starter":0,"pro":0}
    calls = 0
    for k, v in keys.items():
        p = v.get("plan","free")
        if p in plans:
            plans[p] += 1
        calls += usage_cache.get(f"{k}:{today}", 0)
    rev = plans["starter"]*9 + plans["pro"]*29
    return (
        f"👥 **Total Customers:** {len(keys)}\n\n"
        f"🆓 Free: {plans['free']}\n\n"
        f"⭐ Starter ($9): {plans['starter']}\n\n"
        f"🏆 Pro ($29): {plans['pro']}\n\n"
        f"📊 **Total Calls Today:** {calls}\n\n"
        f"💰 **Est. Monthly Revenue:** ${rev}"
    )

def admin_export_json(pw):
    if not _auth(pw):
        return "❌ Wrong password"
    return json.dumps(load_keys(), indent=2)

# ── Build Gradio UI ───────────────────────────────────────────────────────────
gpu_note = "⚡ ZeroGPU active — best mode uses SD1.5 inpainting" if HAS_ZERO_GPU else "🖥️ CPU mode — fast & quality use LaMa. Best = LaMa fallback. Enable ZeroGPU for SD1.5."

with gr.Blocks(title="🪄 FreeME — Magic Eraser API") as demo:

    gr.HTML("""
    <div style="text-align:center;padding:20px 0 8px;">
      <h1 style="font-size:1.8rem;font-weight:800;margin:0;">🪄 FreeME — Magic Eraser API</h1>
      <p style="color:#94a3b8;margin:6px 0 0;">3 Models · API Key Auth · <a href="https://freebg.site" style="color:#f97316;">freebg.site</a></p>
    </div>""")

    with gr.Tabs():

        # ── Tab 1: Magic Erase ────────────────────────────────────────────────
        with gr.TabItem("🪄 Magic Erase"):
            gr.Markdown(f"> {gpu_note}")
            with gr.Row():
                with gr.Column(scale=2):
                    api_key_in  = gr.Textbox(label="🔑 API Key", placeholder="freebg-me-xxxx", type="password")
                    model_in    = gr.Radio(["fast","quality","best"], value="fast", label="⚙️ Model",
                                           info="fast=LaMa | quality=LaMa+sharpen | best=SD1.5 on GPU / LaMa fallback on CPU")
                    img_editor  = gr.ImageEditor(
                        label="🖌️ Upload image then paint over what to erase",
                        type="pil",
                        brush=gr.Brush(colors=["#FF0000"], default_size=20),
                        height=440,
                    )
                    erase_btn   = gr.Button("🪄 Erase Now", variant="primary")
                with gr.Column(scale=1):
                    result_out  = gr.Image(label="✨ Result", type="pil", height=360)
                    status_out  = gr.Textbox(label="Status", interactive=False)

            erase_btn.click(fn=erase_ui, inputs=[img_editor, model_in, api_key_in], outputs=[result_out, status_out])

            gr.HTML("""
            <div style="margin-top:20px;padding:14px 18px;background:#1e293b;border-radius:8px;">
              <b>📦 Plans</b>
              <table style="width:100%;margin-top:8px;border-collapse:collapse;">
                <tr style="border-bottom:1px solid #334155;"><th style="text-align:left;padding:6px;">Plan</th><th>Daily</th><th>Models</th><th>Price</th></tr>
                <tr><td style="padding:5px;">Free</td><td style="text-align:center;">10</td><td style="text-align:center;">fast</td><td style="text-align:center;">$0</td></tr>
                <tr><td style="padding:5px;">Starter</td><td style="text-align:center;">100</td><td style="text-align:center;">fast + quality</td><td style="text-align:center;">$9/mo</td></tr>
                <tr><td style="padding:5px;">Pro</td><td style="text-align:center;">500</td><td style="text-align:center;">all 3</td><td style="text-align:center;">$29/mo</td></tr>
              </table>
              <p style="margin:8px 0 0;color:#94a3b8;font-size:.85rem;">📧 <a href="mailto:support@freebg.site" style="color:#f97316;">support@freebg.site</a></p>
            </div>""")

        # ── Tab 2: My Usage ───────────────────────────────────────────────────
        with gr.TabItem("📊 My Usage"):
            gr.Markdown("### Check Your API Key Usage")
            with gr.Row():
                usage_key_in = gr.Textbox(label="🔑 Your API Key", type="password",
                                          placeholder="freebg-me-xxxx", scale=4)
                usage_btn    = gr.Button("Check Usage", scale=1)
            usage_out = gr.Markdown()
            usage_btn.click(fn=check_usage_ui, inputs=usage_key_in, outputs=usage_out)

        # ── Tab 3: API Docs ───────────────────────────────────────────────────
        with gr.TabItem("📖 API Docs"):
            gr.Markdown(f"""
## FreeME Magic Eraser — REST API

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

**GPU Status:** `{"ZeroGPU Active ⚡" if HAS_ZERO_GPU else "CPU Mode 🖥️"}`

---

### `POST /api/magic-erase`

| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `image` | file | ✅ | Input image JPG/PNG |
| `mask` | file | ✅ | Mask PNG — white=erase, black=keep |
| `model` | string | ❌ | `fast` / `quality` / `best` (default: fast) |

**Header:** `X-API-Key: your-key`

**Response:** `image/png`

---

### `GET /api/usage`

**Header:** `X-API-Key: your-key`

```json
{{
  "plan": "starter",
  "daily_limit": 100,
  "used_today": 23,
  "remaining": 77,
  "models": ["fast", "quality"],
  "owner": "user@email.com"
}}
```

---

### Python
```python
import requests
r = requests.post(
    "https://freebg-magic-eraser.hf.space/api/magic-erase",
    headers={{"X-API-Key": "your-key"}},
    files={{"image": open("photo.jpg","rb"), "mask": open("mask.png","rb")}},
    data={{"model": "quality"}}
)
open("result.png","wb").write(r.content)
```

### JavaScript
```javascript
const fd = new FormData();
fd.append("image", imageFile);
fd.append("mask", maskBlob, "mask.png");
fd.append("model", "fast");
const res = await fetch("https://freebg-magic-eraser.hf.space/api/magic-erase", {{
  method: "POST", headers: {{"X-API-Key": "your-key"}}, body: fd
}});
const url = URL.createObjectURL(await res.blob());
```
""")

        # ── Tab 4: Admin ──────────────────────────────────────────────────────
        with gr.TabItem("🔐 Admin"):
            gr.HTML("<h3 style='margin:12px 0 4px;'>👑 FreeME Admin Panel</h3>")
            
            with gr.Group(visible=True) as login_group:
                admin_pw_login = gr.Textbox(label="Admin Password", type="password", placeholder="Enter admin password")
                login_btn = gr.Button("Login", variant="primary")
                login_err = gr.Markdown(visible=False)
        
            with gr.Group(visible=False) as panel_group:
                admin_pw = gr.State("")   # store password after login

                with gr.Tabs():
    
                    # Add Customer
                    with gr.TabItem("➕ Add Customer"):
                        gr.Markdown("Fill details → Generate Key → Key is **automatically saved** to HF Secret")
                        with gr.Row():
                            ac_email = gr.Textbox(label="Customer Email", placeholder="customer@gmail.com", scale=3)
                            ac_plan  = gr.Dropdown(["free","starter","pro"], value="starter", label="Plan", scale=2)
                            ac_limit = gr.Textbox(label="Custom Daily Limit (optional)", placeholder="leave empty = plan default", scale=2)
                        ac_btn    = gr.Button("🔑 Generate API Key + Save to HF Secret", variant="primary")
                        ac_status = gr.Markdown()
                        ac_key    = gr.Textbox(label="🔑 Generated Key (copy & send to customer)", interactive=False)
                        ac_json   = gr.Code(label="Updated JSON (auto-saved — for manual backup)", language="json")
                        ac_btn.click(fn=admin_add_customer,
                                     inputs=[admin_pw, ac_email, ac_plan, ac_limit],
                                     outputs=[ac_status, ac_key, ac_json])
    
                    # Customers list
                    with gr.TabItem("👥 Customers"):
                        cust_refresh = gr.Button("🔄 Refresh List")
                        cust_out     = gr.Markdown()
                        with gr.Row():
                            upg_key  = gr.Textbox(label="API Key to Edit", placeholder="freebg-me-xx-xxxxxxxx-xxxxxxxx", scale=3)
                            upg_plan = gr.Dropdown(["free","starter","pro","master"], value="pro", label="New Plan", scale=2)
                        upg_btn    = gr.Button("⬆️ Upgrade Plan + Save to HF Secret")
                        upg_status = gr.Textbox(label="Status", interactive=False)
                        upg_json   = gr.Code(label="Updated JSON", language="json")
    
                        gr.Markdown("---")
                        del_key  = gr.Textbox(label="API Key to Delete", placeholder="freebg-me-xx-xxxxxxxx-xxxxxxxx")
                        del_btn  = gr.Button("🗑️ Delete Key + Save to HF Secret", variant="stop")
                        del_status = gr.Textbox(label="Status", interactive=False)
                        del_json   = gr.Code(label="Updated JSON", language="json")
    
                        cust_refresh.click(fn=admin_get_customers, inputs=admin_pw, outputs=cust_out)
                        upg_btn.click(fn=admin_upgrade, inputs=[admin_pw, upg_key, upg_plan], outputs=[upg_status, upg_json])
                        del_btn.click(fn=admin_delete,  inputs=[admin_pw, del_key], outputs=[del_status, del_json])
    
                    # Stats
                    with gr.TabItem("📊 Stats"):
                        stats_btn = gr.Button("🔄 Refresh Stats")
                        stats_out = gr.Markdown()
                        stats_btn.click(fn=admin_stats, inputs=admin_pw, outputs=stats_out)
    
                    # Export JSON
                    with gr.TabItem("📋 Export JSON"):
                        gr.Markdown("Use this if you want to manually paste into HF Secret. **Auto-update already happens** when you Add/Edit/Delete customers.")
                        exp_btn = gr.Button("📋 Get Current JSON")
                        exp_out = gr.Code(language="json")
                        exp_btn.click(fn=admin_export_json, inputs=admin_pw, outputs=exp_out)

                def do_login(pw):
                    if pw == ADMIN_PASSWORD:
                        return gr.update(visible=False), gr.update(visible=True), pw, gr.update(visible=False)
                    return gr.update(visible=True), gr.update(visible=False), "", gr.update(visible=True, value="❌ Wrong password")
                
                login_btn.click(
                    fn=do_login,
                    inputs=admin_pw_login,
                    outputs=[login_group, panel_group, admin_pw, login_err]
                )

# ── FastAPI REST endpoints ────────────────────────────────────────────────────
app = demo.app

@app.post("/api/magic-erase")
async def api_magic_erase(request: Request):
    import traceback
    t0 = time.time()
    try:
        form    = await request.form()
        api_key = request.headers.get("X-API-Key", "")
        model   = form.get("model", "fast").lower()
        if model not in ["fast","quality","best"]:
            model = "fast"

        allowed, err = check_and_count(api_key, model)
        if not allowed:
            return JSONResponse({"success": False, "error": err}, status_code=403)

        img_file  = form.get("image")
        mask_file = form.get("mask")
        if img_file is None or mask_file is None:
            return JSONResponse({"success": False, "error": "Both 'image' and 'mask' files required"}, status_code=400)

        image = Image.open(io.BytesIO(await img_file.read())).convert("RGB")
        mask  = Image.open(io.BytesIO(await mask_file.read())).convert("L")
        if mask.size != image.size:
            mask = mask.resize(image.size, Image.NEAREST)

        fn     = {"fast": run_lama, "quality": run_lama_quality, "best": run_best}[model]
        result = fn(image, mask)

        buf = io.BytesIO()
        result.save(buf, format="PNG")
        return Response(
            content=buf.getvalue(),
            media_type="image/png",
            headers={
                "X-Model":           model,
                "X-Processing-Time": f"{time.time()-t0:.2f}s",
                "X-GPU":             "zerogpu" if HAS_ZERO_GPU else ("cuda" if HAS_CUDA else "cpu"),
            }
        )
    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"success": False, "error": str(e)}, status_code=500)

@app.get("/api/usage")
async def api_usage(request: Request):
    return JSONResponse(get_usage_info(request.headers.get("X-API-Key","")))

@app.get("/api/health")
async def api_health():
    return JSONResponse({
        "status":   "ok",
        "lama":     lama_model is not None,
        "zero_gpu": HAS_ZERO_GPU,
        "cuda":     HAS_CUDA,
        "mode":     "zerogpu" if HAS_ZERO_GPU else ("cuda" if HAS_CUDA else "cpu"),
    })

# ── Launch ────────────────────────────────────────────────────────────────────
load_lama()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False)