"""
سرویس حذف پس‌زمینه — نسخه‌ی اجرا روی Hugging Face Spaces (ZeroGPU، رایگان).

معماری نهایی (با آزمون و خطای واقعی روی همین Space):

  * سخت‌افزار ZeroGPU **خودش** یک سرور Gradio روی پورت پیش‌فرض بالا می‌آورد؛ پس
    اجرای uvicorn خودمان ممکن نیست (`address already in use`) و مسیرهای سفارشی
    هم زیر Gradio گم می‌شوند (۴۰۴).
  * بنابراین ارتباط افزونه از طریق **API خود Gradio** انجام می‌شود و برای اینکه
    هیچ فایل/مسیر موقتی لازم نباشد، ورودی و خروجی **base64** است:

        POST /gradio_api/call/api_cutout   {"data": ["<image b64>", "silueta", "<token>"]}
        → {"event_id": "..."}
        GET  /gradio_api/call/api_cutout/<event_id>
        → {"data": ["<png b64>", "0.5068"]}

  * سخت‌افزار ZeroGPU بدون وجود یک تابع `@spaces.GPU` بالا نمی‌آید
    («No @spaces.GPU function detected during startup»)؛ یک تابع تشریفاتی داریم
    که هرگز صدا زده نمی‌شود و سهمیه‌ی GPU مصرف نمی‌کند.

تله‌های مدل که در کد لحاظ شده:
  * پیش‌پردازش **فقط /255** (بدون نرمال‌سازی ImageNet؛ وگرنه خروجی صفر و ماسک
    خالی می‌شود و کل تصویر پاک می‌گردد).
  * نرمال‌سازی خروجی کمینه-بیشینه (سیگموید ماسک بد می‌دهد).
  * `intra_op_num_threads` محدود، تا روی کانتینر مشترک پایدار بماند.
"""

import base64
import io
import os
import threading
import urllib.request

import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageFilter

try:  # pragma: no cover
    import spaces  # noqa: F401
except Exception:
    spaces = None


if spaces is not None:

    @spaces.GPU(duration=1)
    def _gpu_probe() -> bool:
        """فقط برای الزام ZeroGPU؛ هرگز اجرا نمی‌شود (سهمیه‌ی GPU مصرف نمی‌کند)."""
        return True

else:

    def _gpu_probe() -> bool:
        return True


# کلید مشترک با وردپرس؛ در Space → Settings → Variables and secrets ست می‌شود.
TOKEN = os.environ.get("PCMO_TOKEN", "")

MODELS = {
    # نام: (مسیر فایل، ابعاد ورودی، آدرس دانلود، نرمال‌سازی ImageNet؟)
    #
    # نکته‌ی مهم: مدل‌های خانواده‌ی u2net/isnet **فقط** تقسیم بر ۲۵۵ می‌خواهند،
    # ولی BiRefNet نرمال‌سازی ImageNet هم می‌خواهد (طبق preprocessor_config خودش).
    # اشتباه گرفتن این دو، خروجی صفر و ماسک خالی می‌دهد.
    "silueta": (
        "models/silueta.onnx",
        320,
        "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx",
        False,
    ),
    "u2netp": (
        "models/u2netp.onnx",
        320,
        "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx",
        False,
    ),
    "isnet": (
        "models/isnet-general-use.onnx",
        1024,
        "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx",
        False,
    ),
    # «مدل دقیق» — همان مدلی که رقیب استفاده می‌کند. سنگین است (۴۶۷ مگابایت) و
    # در اولین استفاده دانلود می‌شود؛ برای عکس‌های شلوغ و پس‌زمینه‌ی بافت‌دار
    # (چوب، پارچه) به‌طور محسوس بهتر از isnet است.
    "birefnet": (
        "models/birefnet-fp16.onnx",
        1024,
        "https://huggingface.co/onnx-community/BiRefNet-ONNX/resolve/main/onnx/model_fp16.onnx",
        True,
    ),
}
DEFAULT_MODEL = os.environ.get("PCMO_MODEL", "silueta")

# مدل‌هایی که خروجی‌شان logit است و باید سیگموید شود (BiRefNet).
# بقیه‌ی مدل‌ها (u2net/isnet) خروجی‌شان نسبتاً نرمال است و کمینه-بیشینه کافی است.
SIGMOID_MODELS = {"birefnet"}

MAX_BYTES = int(os.environ.get("PCMO_MAX_MB", "15")) * 1024 * 1024

_lock = threading.Lock()
_sessions: dict[str, tuple[ort.InferenceSession, int]] = {}


def ensure_model(name: str) -> str:
    """فایل مدل را (یک‌بار در عمر کانتینر) دانلود می‌کند."""
    path, _, url, _ = MODELS[name]
    if os.path.exists(path) and os.path.getsize(path) > 1_000_000:
        return path

    os.makedirs("models", exist_ok=True)
    partial = path + ".part"
    print(f"[pcmo] downloading model {name} …", flush=True)
    with urllib.request.urlopen(url, timeout=600) as response, open(partial, "wb") as handle:
        while True:
            chunk = response.read(1024 * 256)
            if not chunk:
                break
            handle.write(chunk)
    os.replace(partial, path)
    print(f"[pcmo] model ready: {os.path.getsize(path) // 1048576} MB", flush=True)
    return path


def session_for(name: str) -> tuple[ort.InferenceSession, int]:
    with _lock:
        if name not in _sessions:
            path = ensure_model(name)
            options = ort.SessionOptions()
            options.intra_op_num_threads = int(os.environ.get("PCMO_THREADS", "2"))
            options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            _sessions[name] = (
                ort.InferenceSession(path, sess_options=options, providers=["CPUExecutionProvider"]),
                MODELS[name][1],
            )
        return _sessions[name]


def build_alpha(image: Image.Image, name: str) -> Image.Image:
    session, size = session_for(name)
    normalize = MODELS[name][3]

    rgb = image.convert("RGB").resize((size, size), Image.LANCZOS)
    pixels = np.asarray(rgb, dtype=np.float32) / 255.0

    # فقط مدل‌های ViT-محور (BiRefNet) نرمال‌سازی ImageNet می‌خواهند.
    if normalize:
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        pixels = (pixels - mean) / std

    tensor = np.transpose(pixels, (2, 0, 1))[None].astype(np.float32)

    raw = session.run(None, {session.get_inputs()[0].name: tensor})[0]
    raw_mask = raw[0, 0] if raw.ndim == 4 else raw[0]

    if name in SIGMOID_MODELS:
        # خروجی این مدل‌ها logit است (مثلاً ‎-18‎ تا ‎+5‎) و باید سیگموید شود.
        # با کمینه-بیشینه، همه‌ی مقادیر نزدیک هم به «سوژه» تبدیل می‌شدند و پوشش
        # به ۹۹.۹۵٪ می‌رسید (کل تصویر دست‌نخورده می‌ماند).
        mask = 1.0 / (1.0 + np.exp(-np.clip(raw_mask.astype(np.float32), -30.0, 30.0)))
    else:
        span = float(raw_mask.max() - raw_mask.min())
        mask = (raw_mask - raw_mask.min()) / (span if span > 1e-6 else 1.0)

    # ۱) آستانه‌ی نرم.
    #
    # ماسک خام مدل، مقادیر ضعیفی (۰.۲ تا ۰.۵) در اطراف سوژه دارد که بعد از بزرگ
    # شدن به شکل «هاله‌ی رنگی» و «لکه‌های نیمه‌شفاف پس‌زمینه» دیده می‌شود. با یک
    # باند نرم، این مقادیر صفر می‌شوند و لبه‌ی سوژه تمیز ولی نرم می‌ماند.
    threshold = float(os.environ.get("PCMO_MASK_THRESHOLD", "0.45"))
    band = max(1e-6, float(os.environ.get("PCMO_MASK_BAND", "0.12")))
    mask = np.clip((mask - threshold) / band, 0.0, 1.0)

    alpha = Image.fromarray((mask * 255).astype(np.uint8), mode="L")

    # ۲) پاک‌سازی نقطه‌های ریز: «باز کردن» (فرسایش سپس گسترش) لکه‌های کوچک را
    #    حذف می‌کند و به برگ و شاخه‌ی نازک کاری ندارد (کرنل کوچک انتخاب شد).
    alpha = alpha.filter(ImageFilter.MinFilter(3)).filter(ImageFilter.MaxFilter(3))

    # ۳) بازگشت به ابعاد تصویر با درون‌یابی نرم.
    #
    # LANCZOS روی ماسک، لبه‌های پله‌ای و «رینگ» رنگی می‌سازد؛ BILINEAR تمیزتر است.
    return alpha.resize(image.size, Image.BILINEAR)


def cut_out(image: Image.Image, name: str, feather: int = 0) -> tuple[Image.Image, float]:
    alpha = build_alpha(image, name)

    # کمی نرمی روی لبه (برای چاپ DTF بهتر است) + نرمی اختیاری بیشتر از پارامتر.
    blur = max(0.6, float(feather))
    alpha = alpha.filter(ImageFilter.GaussianBlur(blur))

    output = image.convert("RGBA")
    output.putalpha(alpha)
    coverage = float((np.asarray(alpha) > 128).mean())
    return output, coverage


def ui_cutout(image, model_name):
    """رابط کاربری مرورگر (برای تست دستی)."""
    if image is None:
        return None
    output, _ = cut_out(image, model_name or DEFAULT_MODEL, 0)
    return output


def api_cutout(image_b64: str, model_name: str, token: str):
    """
    اندپوینت افزونه‌ی وردپرس: تصویر base64 می‌گیرد و PNG شفاف base64 + پوشش
    ماسک برمی‌گرداند. توکن برای جلوگیری از استفاده‌ی ناخواسته بررسی می‌شود.
    """
    if TOKEN and (token or "").strip() != TOKEN:
        return "", "ERR:token"

    name = (model_name or DEFAULT_MODEL).strip()
    if name not in MODELS:
        return "", "ERR:model"

    payload = (image_b64 or "").strip()
    if payload.startswith("data:"):
        payload = payload.split(",", 1)[-1]
    if not payload:
        return "", "ERR:empty"

    try:
        raw_bytes = base64.b64decode(payload, validate=False)
    except Exception:
        return "", "ERR:decode"

    if len(raw_bytes) > MAX_BYTES:
        return "", "ERR:big"

    try:
        image = Image.open(io.BytesIO(raw_bytes))
        image.load()
    except Exception:
        return "", "ERR:image"

    output, coverage = cut_out(image, name, 0)
    buffer = io.BytesIO()
    output.save(buffer, format="PNG", optimize=True)

    return base64.b64encode(buffer.getvalue()).decode("ascii"), f"{coverage:.4f}"


def probe_model(model_name: str) -> str:
    """
    اطلاعات تشخیصی مدل: نام/شکل/نوع ورودی‌ها و خروجی‌ها + آمار خروجی خام.

    چرا لازم شد: BiRefNet خروجی‌اش با isnet/u2net فرق دارد و بدون دیدن آمار
    واقعی (کمینه/بیشینه/میانه) نمی‌توان تصمیم گرفت سیگموید لازم است یا نه.
    """
    import json as _json

    name = (model_name or DEFAULT_MODEL).strip()
    if name not in MODELS:
        return "ERR:model"

    session, size = session_for(name)
    info = {
        "model": name,
        "normalize": bool(MODELS[name][3]),
        "size": size,
        "inputs": [{"name": i.name, "shape": list(i.shape), "type": i.type} for i in session.get_inputs()],
        "outputs": [{"name": o.name, "shape": list(o.shape), "type": o.type} for o in session.get_outputs()],
    }

    pixels = np.full((size, size, 3), 0.5, dtype=np.float32)
    if MODELS[name][3]:
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        pixels = (pixels - mean) / std

    tensor = np.transpose(pixels, (2, 0, 1))[None].astype(np.float32)
    outputs = session.run(None, {session.get_inputs()[0].name: tensor})

    stats = []
    for index, raw in enumerate(outputs):
        flat = np.asarray(raw).reshape(-1).astype(np.float32)
        stats.append(
            {
                "index": index,
                "dims": list(np.asarray(raw).shape),
                "min": round(float(flat.min()), 4),
                "max": round(float(flat.max()), 4),
                "mean": round(float(flat.mean()), 4),
                "p50": round(float(np.percentile(flat, 50)), 4),
                "p90": round(float(np.percentile(flat, 90)), 4),
            }
        )

    info["stats"] = stats
    return _json.dumps(info, ensure_ascii=False)


with gr.Blocks(title="حذف پس‌زمینه") as demo:
    gr.Markdown("# حذف پس‌زمینه\nسرویس افزونه‌ی چاپ سفارشی وردپرس.")

    with gr.Row():
        ui_image = gr.Image(type="pil", label="تصویر")
        ui_out = gr.Image(type="pil", label="بدون پس‌زمینه")

    ui_model = gr.Dropdown(choices=list(MODELS), value=DEFAULT_MODEL, label="مدل")
    ui_btn = gr.Button("حذف پس‌زمینه", variant="primary")
    ui_btn.click(ui_cutout, inputs=[ui_image, ui_model], outputs=[ui_out], api_name="ui_cutout")

    gr.Markdown("---\n### اندپوینت API افزونه (base64)")
    api_in = gr.Textbox(label="ورودی base64")
    api_model = gr.Dropdown(choices=list(MODELS), value=DEFAULT_MODEL, label="مدل")
    api_token = gr.Textbox(label="توکن")
    api_btn = gr.Button("اجرا")
    api_out = gr.Textbox(label="خروجی base64")
    api_cov = gr.Textbox(label="پوشش ماسک")
    api_btn.click(api_cutout, inputs=[api_in, api_model, api_token], outputs=[api_out, api_cov], api_name="api_cutout")
    probe_btn = gr.Button("اطلاعات مدل (تشخیصی)")
    probe_btn.click(probe_model, inputs=[api_model], outputs=[api_out], api_name="probe_model")


def _warm() -> None:
    """
    آماده‌سازی مدل‌ها پیش از اولین درخواست کاربر.

    `isnet` (۱۷۰ مگابایت) در اولین استفاده دانلود می‌شود؛ اگر همین‌جا و در پس‌زمینه
    انجام شود، اولین کلیک مشتری معطل دانلود نمی‌ماند.
    """
    for name in ("silueta", "isnet", "birefnet"):
        try:
            session_for(name)
            print(f"[pcmo] warm: {name} ready", flush=True)
        except Exception as error:  # noqa: BLE001
            print(f"[pcmo] warm-up failed for {name}: {error}", flush=True)


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


if __name__ == "__main__":
    demo.launch()