from fastapi import FastAPI
from pydantic import BaseModel
from transparent_background import Remover
from PIL import Image
import io
import requests
import base64
import warnings

warnings.filterwarnings("ignore")

app = FastAPI(title="Background Remover API")

# Chargé une seule fois
remover = Remover(mode="base-nightly")


class RemoveBackgroundRequest(BaseModel):
    image_url: str
    sod_type: str = "rgba"


class RemoveBackgroundResponse(BaseModel):
    image_base64: str


@app.get("/")
def welcome():
    return {"message": "welcome"}


@app.post("/remove-background", response_model=RemoveBackgroundResponse)
def remove_background(payload: RemoveBackgroundRequest):
    # Téléchargement de l'image depuis l'URL
    response = requests.get(payload.image_url, timeout=10)
    response.raise_for_status()
    pil_image = Image.open(io.BytesIO(response.content)).convert("RGB")

    # Traitement du fond
    result = remover.process(pil_image, type=payload.sod_type)

    # Conversion en PNG en mémoire
    buf = io.BytesIO()
    result.save(buf, format="PNG")
    buf.seek(0)

    # Encodage en Base64
    img_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")

    return {"image_base64": img_base64}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)