from __future__ import annotations

import os
import threading
from dataclasses import dataclass

import numpy as np
from PIL import Image

try:
    import spaces

    gpu_task = spaces.GPU(duration=180)
except ImportError:
    def gpu_task(function):
        return function


MODEL_ID = os.getenv("MODEL_ID", "ZhengPeng7/BiRefNet")
MODEL_REVISION = os.getenv(
    "MODEL_REVISION", "e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4"
)
MODEL_INPUT_SIZE = int(os.getenv("MODEL_INPUT_SIZE", "1024"))

_runtime: "BiRefNetRuntime | None" = None
_load_lock = threading.Lock()


@dataclass
class BiRefNetRuntime:
    model: object
    device: object
    dtype: object

    @classmethod
    def load(cls) -> "BiRefNetRuntime":
        import torch
        from transformers import AutoModelForImageSegmentation

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        dtype = torch.float16 if device.type == "cuda" else torch.float32
        torch.set_float32_matmul_precision("high")
        model = AutoModelForImageSegmentation.from_pretrained(
            MODEL_ID,
            revision=MODEL_REVISION,
            trust_remote_code=True,
        )
        model.to(device=device, dtype=dtype)
        model.eval()
        return cls(model=model, device=device, dtype=dtype)

    def predict(self, image: Image.Image) -> Image.Image:
        import torch
        from torchvision.transforms import functional as TF

        boxed, content_box = _letterbox(image, MODEL_INPUT_SIZE)
        tensor = TF.to_tensor(boxed)
        tensor = TF.normalize(
            tensor,
            mean=(0.485, 0.456, 0.406),
            std=(0.229, 0.224, 0.225),
        ).unsqueeze(0)
        tensor = tensor.to(device=self.device, dtype=self.dtype)

        with torch.inference_mode():
            prediction = self.model(tensor)
            logits = _last_tensor(prediction)
            probability = logits.sigmoid()[0].squeeze().float().cpu().numpy()

        probability = np.clip(probability * 255.0, 0, 255).astype(np.uint8)
        square_mask = Image.fromarray(probability, mode="L")
        content_mask = square_mask.crop(content_box)
        return content_mask.resize(image.size, Image.Resampling.LANCZOS)


def _letterbox(image: Image.Image, size: int) -> tuple[Image.Image, tuple[int, int, int, int]]:
    scale = min(size / image.width, size / image.height)
    resized_size = (max(1, round(image.width * scale)), max(1, round(image.height * scale)))
    resized = image.resize(resized_size, Image.Resampling.LANCZOS)
    # ImageNet mean becomes approximately zero after normalization.
    canvas = Image.new("RGB", (size, size), (124, 116, 104))
    left = (size - resized.width) // 2
    top = (size - resized.height) // 2
    canvas.paste(resized, (left, top))
    return canvas, (left, top, left + resized.width, top + resized.height)


def _last_tensor(value):
    import torch

    if torch.is_tensor(value):
        return value
    if hasattr(value, "logits"):
        return _last_tensor(value.logits)
    if isinstance(value, (tuple, list)) and value:
        for item in reversed(value):
            try:
                return _last_tensor(item)
            except (TypeError, ValueError):
                continue
    raise TypeError("BiRefNet beklenmeyen bir çıktı biçimi döndürdü.")


def get_runtime() -> BiRefNetRuntime:
    global _runtime
    if _runtime is None:
        with _load_lock:
            if _runtime is None:
                _runtime = BiRefNetRuntime.load()
    return _runtime


@gpu_task
def predict_mask(image: Image.Image) -> Image.Image:
    """Run segmentation on ZeroGPU/classical GPU, with transparent CPU fallback."""
    return get_runtime().predict(image)