import torch
from transformers import AutoModelForImageSegmentation
from torchvision import transforms
from PIL import Image
import numpy as np

# Device setup
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Global models
_birefnet = None
_birefnet_lite = None

# Preprocessing transform
_transform = transforms.Compose(
    [
        transforms.Resize((768, 768)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


def load_model(fast_mode=False):
    global _birefnet, _birefnet_lite

    if fast_mode:
        if _birefnet_lite is None:
            _birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
                "ZhengPeng7/BiRefNet_lite", trust_remote_code=True
            ).to(_device)
        return _birefnet_lite

    else:
        if _birefnet is None:
            _birefnet = AutoModelForImageSegmentation.from_pretrained(
                "ZhengPeng7/BiRefNet", trust_remote_code=True
            ).to(_device)
        return _birefnet


def unload_model():
    global _birefnet, _birefnet_lite
    _birefnet = None
    _birefnet_lite = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _apply_mask(image: Image.Image, mask_tensor, bg, bg_type: str):
    size = image.size
    mask = transforms.ToPILImage()(mask_tensor.squeeze()).resize(size)

    if bg_type == "Color" and isinstance(bg, str) and bg.startswith("#"):
        color_rgb = tuple(int(bg[i : i + 2], 16) for i in (1, 3, 5))
        bg_image = Image.new("RGBA", size, color_rgb + (255,))
    elif bg_type in ["Image", "Video"]:
        bg_image = bg.convert("RGBA").resize(size)
    else:
        raise ValueError("Invalid background type or format")

    return Image.composite(image.convert("RGBA"), bg_image, mask)


def _get_transform(image_size: tuple[int, int] | None = None):
    if image_size:
        resize = transforms.Resize(image_size)
    else:
        resize = transforms.Resize((768, 768))  # default fallback

    return transforms.Compose(
        [
            resize,
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )


def predict(
    input: Image.Image | list[Image.Image],
    bg: str | Image.Image,
    bg_type: str = "Color",
    fast_mode: bool = True,
    bg_frames: list[np.ndarray] | None = None,
    image_size_mode: str = "default",  # "default", "origin", or "custom"
    custom_size: tuple[int, int] | None = None,
) -> Image.Image | list[Image.Image]:

    model = _birefnet_lite if fast_mode else _birefnet
    if model is None:
        raise RuntimeError("Model not loaded. Please call `load_model()` first.")

    def get_transform_for_img(img: Image.Image):
        if image_size_mode == "origin":
            return _get_transform(img.size)
        elif image_size_mode == "custom":
            if not custom_size:
                raise ValueError("`custom_size` must be provided for 'custom' mode.")
            return _get_transform(custom_size)
        else:
            return _get_transform()  # default to (768, 768)

    if isinstance(input, Image.Image):
        _transform = get_transform_for_img(input)
        input_tensor = _transform(input).unsqueeze(0).to(_device)
        with torch.no_grad():
            pred = model(input_tensor)[-1].sigmoid().cpu()
        bg_image = (
            Image.fromarray(bg_frames[0]) if bg_type == "Video" and bg_frames else bg
        )
        return _apply_mask(input, pred[0], bg_image, bg_type)

    elif isinstance(input, list):
        results = []
        for idx, frame in enumerate(input):
            _transform = get_transform_for_img(frame)
            input_tensor = _transform(frame).unsqueeze(0).to(_device)
            with torch.no_grad():
                pred = model(input_tensor)[-1].sigmoid().cpu()
            if bg_type == "Video" and bg_frames:
                bg_image = Image.fromarray(bg_frames[idx % len(bg_frames)])
            else:
                bg_image = bg
            results.append(_apply_mask(frame, pred[0], bg_image, bg_type))
        return results

    else:
        raise TypeError("`input` must be a PIL.Image or a list of PIL.Image.")


# # Default behavior (768x768)
# predict(img, "#ffffff", "Color")

# # Keep original image size
# predict(img, "#ffffff", "Color", image_size_mode="origin")

# # Resize to 1024x1024
# predict(img, "#ffffff", "Color", image_size_mode="custom", custom_size=(1024, 1024))