"""
Toolmingo Background Remover — besplatni API (Hugging Face Docker Space).
BiRefNet_lite (MIT) preko onnxruntime. POST /remove -> providni PNG.
"""
import io
import numpy as np
import onnxruntime as ort
from PIL import Image
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download

MODEL_REPO = "onnx-community/BiRefNet_lite-ONNX"
MODEL_FILE = "onnx/model.onnx"  # fp32, najcistije ivice
SIZE = 1024
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(3, 1, 1)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(3, 1, 1)

print("Skidam model...", flush=True)
_model_path = hf_hub_download(MODEL_REPO, MODEL_FILE)
_sess = ort.InferenceSession(_model_path, providers=["CPUExecutionProvider"])
_inp = _sess.get_inputs()[0].name
print("Model spreman.", flush=True)

app = FastAPI(title="Toolmingo Background Remover")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def cut(img: Image.Image) -> Image.Image:
    img = img.convert("RGB")
    w, h = img.size
    small = img.resize((SIZE, SIZE), Image.BILINEAR)
    x = np.asarray(small, dtype=np.float32) / 255.0
    x = x.transpose(2, 0, 1)
    x = (x - MEAN) / STD
    x = x[None, ...]
    out = _sess.run(None, {_inp: x})[0]  # [1,1,1024,1024]
    mask = 1.0 / (1.0 + np.exp(-out[0, 0]))  # sigmoid
    mask = (mask * 255).astype(np.uint8)
    mask_img = Image.fromarray(mask).resize((w, h), Image.BILINEAR)
    res = img.convert("RGBA")
    res.putalpha(mask_img)
    return res


@app.get("/")
def root():
    return {"status": "ok", "model": MODEL_REPO}


@app.post("/remove")
async def remove(file: UploadFile = File(...)):
    try:
        data = await file.read()
        img = Image.open(io.BytesIO(data))
        res = cut(img)
        buf = io.BytesIO()
        res.save(buf, format="PNG")
        return Response(content=buf.getvalue(), media_type="image/png")
    except Exception as e:
        return JSONResponse(status_code=400, content={"error": str(e)})