server.py

import os
import io
import json
import requests
import tempfile
from datetime import date
from fastapi import FastAPI, UploadFile, File, HTTPException, Header
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from gradio_client import Client, handle_file

# --- KONFIGURASI (Silakan sesuaikan) ---

# 1. API Key dari remove.bg (dari screenshot Anda)
#    SANGAT DISARANKAN untuk mengatur ini sebagai environment variable di platform deploy Anda.
REMOVEBG_API_KEY = os.getenv("REMOVEBG_API_KEY")

# 2. Nama Gradio Space untuk fallback
BRIA_SPACE = "briaai/BRIA-RMBG-2.0" # Ganti jika Anda punya nama Space Bria yang lain
MYSTYC_SPACE = "Mystyc/remove-bg-rmbg" # Space yang Anda gunakan sekarang

# 3. Pengaturan Jatah Harian untuk remove.bg
DAILY_QUOTA_LIMIT = 2
QUOTA_FILE = os.path.join(tempfile.gettempdir(), "quota.json")

# --- Inisialisasi Aplikasi FastAPI ---
app = FastAPI(title="Dynamic Fallback RMBG Proxy")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# --- FUNGSI-FUNGSI PEMBANTU ---

def _check_and_update_quota():
    """Membaca file quota, memeriksa apakah jatah masih ada, dan mengupdate jika perlu."""
    today = str(date.today())
    quota_data = {"date": today, "count": 0}

    if os.path.exists(QUOTA_FILE):
        try:
            with open(QUOTA_FILE, "r") as f:
                data = json.load(f)
            # Jika tanggal sudah berbeda, reset jatah
            if data.get("date") == today:
                quota_data = data
            else:
                print(f"Tanggal baru terdeteksi. Mereset jatah harian untuk {today}.")
        except (json.JSONDecodeError, FileNotFoundError):
            print("File quota rusak atau tidak ditemukan, membuat yang baru.")
    
    with open(QUOTA_FILE, "w") as f:
        json.dump(quota_data, f)

    if quota_data["count"] < DAILY_QUOTA_LIMIT:
        return True
    else:
        print(f"Jatah harian remove.bg ({DAILY_QUOTA_LIMIT}) untuk tanggal {today} sudah habis.")
        return False

def _increment_quota_count():
    """Menambah hitungan jatah setelah pemanggilan API remove.bg berhasil."""
    try:
        with open(QUOTA_FILE, "r+") as f:
            data = json.load(f)
            data["count"] += 1
            f.seek(0)
            json.dump(data, f)
            f.truncate()
        print(f"Jatah remove.bg digunakan: {data['count']}/{DAILY_QUOTA_LIMIT}")
    except Exception as e:
        print(f"Gagal menambah hitungan jatah: {e}")

def _predict_from_removebg_api(image_blob: bytes):
    """Memanggil API resmi remove.bg."""
    print("Mencoba Prioritas 1: remove.bg API...")
    response = requests.post(
        'https://api.remove.bg/v1/removebg',
        files={'image_file': image_blob},
        headers={'X-Api-Key': REMOVEBG_API_KEY},
        timeout=180
    )
    response.raise_for_status() # Akan error jika status bukan 2xx
    _increment_quota_count() # Tambah hitungan jika berhasil
    print("✅ Sukses menggunakan remove.bg API.")
    return StreamingResponse(io.BytesIO(response.content), media_type="image/png")

# GANTI SELURUH FUNGSI LAMA DENGAN YANG INI

def _predict_from_gradio_space(local_path: str, space_name: str, api_name: str):
    """Fungsi yang direfaktor untuk memanggil Gradio Space manapun dengan api_name dinamis."""
    print(f"Mencoba Prioritas Fallback: Gradio Space '{space_name}' dengan API '{api_name}'...")
    client = Client(src=space_name)
    result = client.predict(image=handle_file(local_path), api_name=api_name)
    
    # --- LOGIKA BARU UNTUK MEMBACA HASIL ---
    image_path = None
    # Jika hasilnya adalah list atau tuple (seperti di Mystyc Space yang baru)
    if isinstance(result, (list, tuple)) and len(result) > 0:
        # Ambil elemen PERTAMA (indeks 0) sebagai path gambar
        image_path = result[0]
    # Jika hasilnya hanya string (seperti di Gradio Space yang lama)
    elif isinstance(result, str):
        image_path = result

    # Lanjutkan proses jika path gambar valid
    if image_path and isinstance(image_path, str) and os.path.exists(image_path):
        with open(image_path, "rb") as f:
            print(f"✅ Sukses menggunakan Gradio Space '{space_name}'.")
            return StreamingResponse(io.BytesIO(f.read()), media_type="image/png")
    
    # Jika format tidak dikenali sama sekali
    raise ValueError(f"Output dari Gradio Space '{space_name}' tidak dikenali: {repr(result)}")

# --- ENDPOINTS ---

@app.get("/health")
def health():
    return {"ok": True, "services": {"primary": "remove.bg", "fallback_1": BRIA_SPACE, "fallback_2": MYSTYC_SPACE}}

@app.post("/remove_bg", response_class=StreamingResponse)
async def remove_bg(file: UploadFile = File(...)):
    blob = await file.read()
    if not blob:
        raise HTTPException(400, "File kosong")

    # --- LOGIKA FALLBACK DIMULAI DI SINI ---

    # 1. Coba remove.bg API jika jatah masih ada
    if REMOVEBG_API_KEY != "MASUKKAN_API_KEY_REMOVEBG_ANDA_DI_SINI" and _check_and_update_quota():
        try:
            return _predict_from_removebg_api(blob)
        except Exception as e:
            print(f"⚠️ Gagal dengan remove.bg API: {e}. Mencoba fallback berikutnya.")

    # Jika remove.bg gagal atau jatah habis, buat file sementara untuk Gradio
    suffix = os.path.splitext(file.filename or ".png")[1]
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(blob)
            tmp_path = tmp.name

        # 2. Coba Bria RM-BG
        try:
            return _predict_from_gradio_space(tmp_path, BRIA_SPACE, api_name="/image")
        except Exception as e:
            print(f"⚠️ Gagal dengan Gradio Space '{BRIA_SPACE}': {e}. Mencoba fallback terakhir.")
        
        # 3. Coba Mystyc RM-BG
        try:
            return _predict_from_gradio_space(tmp_path, MYSTYC_SPACE, api_name="/predict")
        except Exception as e:
            print(f"⚠️ Gagal dengan Gradio Space '{MYSTYC_SPACE}': {e}. Semua layanan gagal.")
            raise HTTPException(status_code=503, detail=f"Semua layanan penghapus background tidak tersedia. Error terakhir: {e}")

    finally:
        # Selalu hapus file sementara
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)