# /app/server/ormbg.py
from __future__ import annotations
import io
import os
from typing import Optional, Union

from rembg import remove as rembg_remove, new_session
from PIL import Image
import onnxruntime as ort

# Choose GPU if available, else CPU
def _default_providers() -> list[str]:
    providers = ort.get_available_providers()
    return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in providers else ["CPUExecutionProvider"]

class ORMBGProcessor:
    """
    Minimal adapter used by server.py.
    Provides:
      - .warmup()
      - .process_bytes(input_bytes) -> bytes (PNG with alpha)
      - .process_pil(img: PIL.Image.Image) -> PIL.Image.Image
      - .close()  (no-op; for API symmetry)
    """
    def __init__(self, model_name: str = "isnet-general-use", providers: Optional[list[str]] = None):
        # Some rembg internals pick providers automatically; onnxruntime-gpu presence
        # will make CUDA provider available. We still compute a list so we can log/verify.
        self.providers = providers or _default_providers()
        # rembg's API chooses the best available EP under the hood if ORT-GPU is present.
        # Just ensure the model name exists; common: isnet-general-use, u2net, u2netp, u2net_human_seg
        self.model_name = model_name
        self.session = new_session(self.model_name)

    def warmup(self) -> None:
        # Tiny 1x1 transparent PNG to trigger model load
        img = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
        _ = self.process_pil(img)

    def process_bytes(self, data: bytes, return_png: bool = True) -> bytes:
        """
        Accept raw image bytes (jpg/png/etc), return PNG bytes with alpha.
        """
        out = rembg_remove(data, session=self.session)
        if return_png:
            # rembg already returns PNG bytes by default; if it's bytes, just pass through
            if isinstance(out, (bytes, bytearray)):
                return bytes(out)
            # if it's a PIL image, convert to PNG bytes
            if isinstance(out, Image.Image):
                buf = io.BytesIO()
                out.save(buf, format="PNG")
                return buf.getvalue()
        # fallback: ensure bytes
        if isinstance(out, Image.Image):
            buf = io.BytesIO()
            out.save(buf, format="PNG")
            return buf.getvalue()
        return bytes(out)

    def process_pil(self, img: Image.Image) -> Image.Image:
        out = rembg_remove(img, session=self.session)
        if isinstance(out, Image.Image):
            return out
        # if rembg returned bytes, decode back to PIL
        return Image.open(io.BytesIO(out)).convert("RGBA")

    def close(self) -> None:
        # Nothing to dispose explicitly; keep for API compatibility
        pass