import os
import sys
import pathlib
import traceback
import urllib.request

import gradio as gr
import onnxruntime as ort
from huggingface_hub import snapshot_download

# =========================
# CONFIG — use secrets only
# =========================
PRIVATE_SPACE_ID     = os.getenv("PRIVATE_SPACE_ID")      
PRIVATE_PKG_IMPORT   = os.getenv("PRIVATE_PKG_IMPORT")    
PRIVATE_MODEL_REL    = os.getenv("PRIVATE_MODEL_REL")     
HF_TOKEN             = os.getenv("HF_TOKEN")              

if not (PRIVATE_SPACE_ID and PRIVATE_PKG_IMPORT and PRIVATE_MODEL_REL and HF_TOKEN):
    raise RuntimeError("Missing one or more required secrets")

# Public fallback
PUBLIC_MODEL_PATH = pathlib.Path(os.getenv("PUBLIC_MODEL_PATH"))
PUBLIC_MODEL_URL  = os.getenv("PUBLIC_MODEL_URL")

if not (PUBLIC_MODEL_PATH and PUBLIC_MODEL_URL):
    raise RuntimeError("Missing PUBLIC_MODEL_PATH or PUBLIC_MODEL_URL in secrets.")

# =========================
# DOWNLOAD PRIVATE SNAPSHOT
# =========================
CACHE_DIR = "private_code"
local_dir = snapshot_download(
    repo_id=PRIVATE_SPACE_ID,
    repo_type="space",
    token=HF_TOKEN,
    local_dir=CACHE_DIR,
    local_dir_use_symlinks=False,
)

sys.path.insert(0, local_dir)

# Import private function
try:
    mod = __import__(PRIVATE_PKG_IMPORT, fromlist=["remove_background_pil"])
    remove_background_pil = getattr(mod, "remove_background_pil")
except Exception:
    raise RuntimeError("Failed to import private function. Ensure PRIVATE_PKG_IMPORT is correct.")

# =========================
# RESOLVE MODEL PATH
# =========================
def resolve_model_path() -> str:
    private_model = pathlib.Path(local_dir) / PRIVATE_MODEL_REL
    if private_model.exists() and private_model.stat().st_size > 5_000_000:
        return str(private_model)

    PUBLIC_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    if not PUBLIC_MODEL_PATH.exists() or PUBLIC_MODEL_PATH.stat().st_size <= 5_000_000:
        urllib.request.urlretrieve(PUBLIC_MODEL_URL, PUBLIC_MODEL_PATH)
    return str(PUBLIC_MODEL_PATH)

MODEL_PATH = resolve_model_path()

# =========================
# HEALTH CHECK
# =========================
def health_check() -> str:
    try:
        exists = os.path.exists(MODEL_PATH)
        size = os.path.getsize(MODEL_PATH) if exists else 0
        providers = ort.get_available_providers()
        return f"Model OK: {exists}, size={size}, providers={providers}"
    except Exception as e:
        return f"Health check failed: {e}"

# =========================
# INFERENCE
# =========================
def infer(image, refine=True):
    if image is None:
        return None
    try:
        return remove_background_pil(image, model_path=MODEL_PATH, refine=refine)
    except Exception:
        raise gr.Error("Background removal failed. Please retry.")

# =========================
# UI
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🪄 AI Background Remover (Public Demo)")
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Upload image")
            refine = gr.Checkbox(value=True, label="Refine edges")
            run_btn = gr.Button("Remove Background", variant="primary")
            diag_btn = gr.Button("Run Health Check")
        with gr.Column():
            out = gr.Image(type="pil", label="Result (PNG with alpha)")
            diag = gr.Code(label="Diagnostics", interactive=False)

    run_btn.click(fn=infer, inputs=[inp, refine], outputs=out)
    diag_btn.click(fn=health_check, inputs=None, outputs=diag)

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