"""RMBG 2.0 model loading and mask generation."""

from pathlib import Path
from typing import Any

from PIL import Image, ImageOps, UnidentifiedImageError


class RmbgError(Exception):
    """Report an expected model, image, device, or output error."""


def validateModelPath(model_path: Path) -> Path:
    """Validate and resolve a local Hugging Face model directory."""
    if not model_path.exists():
        raise RmbgError(f"model directory does not exist: {model_path}")
    if not model_path.is_dir():
        raise RmbgError(f"model path is not a directory: {model_path}")

    config_path = model_path / "config.json"
    if not config_path.is_file():
        raise RmbgError(
            f"model directory does not contain config.json: {model_path}"
        )
    return model_path.resolve()


def selectDevice(torch_module: Any, device_name: str) -> str:
    """Select an available PyTorch inference device."""
    if device_name == "auto":
        if torch_module.cuda.is_available():
            return "cuda"
        mps = getattr(torch_module.backends, "mps", None)
        if mps is not None and mps.is_available():
            return "mps"
        return "cpu"

    if device_name == "cuda" and not torch_module.cuda.is_available():
        raise RmbgError("CUDA was requested but is not available")
    if device_name == "mps":
        mps = getattr(torch_module.backends, "mps", None)
        if mps is None or not mps.is_available():
            raise RmbgError("MPS was requested but is not available")
    return device_name


def loadModel(model_path: Path, device: str) -> Any:
    """Load RMBG from a validated local directory without network access."""
    from transformers import AutoModelForImageSegmentation

    try:
        model = AutoModelForImageSegmentation.from_pretrained(
            str(model_path),
            local_files_only=True,
            trust_remote_code=True,
        )
    except Exception as error:
        raise RmbgError(f"could not load local model: {error}") from error

    return model.eval().to(device)


def openImage(input_path: Path) -> Image.Image:
    """Open an image, apply EXIF orientation, and convert it to RGB."""
    try:
        with Image.open(input_path) as source_image:
            return ImageOps.exif_transpose(source_image).convert("RGB")
    except FileNotFoundError as error:
        raise RmbgError(f"input image does not exist: {input_path}") from error
    except (OSError, UnidentifiedImageError) as error:
        raise RmbgError(f"could not read input image: {error}") from error


def inferMask(
    model: Any,
    image: Image.Image,
    processing_resolution: int,
    device: str,
) -> Image.Image:
    """Infer a soft grayscale mask and restore the source dimensions."""
    import torch
    from torchvision import transforms

    transform_image = transforms.Compose(
        [
            transforms.Resize(
                (processing_resolution, processing_resolution),
                antialias=True,
            ),
            transforms.ToTensor(),
            transforms.Normalize(
                [0.485, 0.456, 0.406],
                [0.229, 0.224, 0.225],
            ),
        ]
    )
    input_tensor = transform_image(image).unsqueeze(0).to(device)

    try:
        with torch.inference_mode():
            prediction = model(input_tensor)[-1].sigmoid().cpu()
    except RuntimeError as error:
        raise RmbgError(f"model inference failed: {error}") from error

    mask_tensor = prediction[0].squeeze().clamp(0, 1)
    processing_mask = transforms.ToPILImage()(mask_tensor).convert("L")
    return processing_mask.resize(image.size, Image.Resampling.BILINEAR)


def saveMask(mask: Image.Image, output_path: Path) -> None:
    """Save an 8-bit grayscale mask to the requested path."""
    try:
        mask.save(output_path)
    except (OSError, ValueError) as error:
        raise RmbgError(f"could not save output mask: {error}") from error


def generateMask(
    model_path: Path,
    input_path: Path,
    output_path: Path,
    processing_resolution: int = 1024,
    device_name: str = "auto",
) -> None:
    """Generate and save a soft foreground mask with a local RMBG model."""
    if processing_resolution <= 0:
        raise RmbgError("processing resolution must be positive")

    local_model_path = validateModelPath(model_path)

    import torch

    device = selectDevice(torch, device_name)
    image = openImage(input_path)
    model = loadModel(local_model_path, device)
    mask = inferMask(model, image, processing_resolution, device)
    saveMask(mask, output_path)