from __future__ import annotations

import numpy as np
from PIL import Image

from .settings import SegmentationSettings


class BiRefNetSegmenter:
    """Semantic subject segmentation via BiRefNet.

    Produces a soft foreground mask that defines subject *topology*: it includes
    colour-contaminated regions (e.g. green-tinted hair) as foreground and
    excludes see-through holes (e.g. gaps between fingers) as background -- the
    distinction that chroma/colour cues cannot make. The mask drives trimap
    construction; ViTMatte still refines the soft edges.

    Environment notes:
    - Weights load from HuggingFace. Behind a firewall set ``HF_ENDPOINT``
      (e.g. ``https://hf-mirror.com``); a flaky proxy may need to be bypassed.
    - The repo ships custom modeling code, so ``trust_remote_code=True`` is
      required (it executes that code -- only use a model source you trust).
    """

    def __init__(self, settings: SegmentationSettings):
        try:
            import torch
            from torchvision import transforms
            from transformers import AutoModelForImageSegmentation
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "Missing segmentation dependencies. Install them with: "
                "pip install torch torchvision transformers timm einops kornia"
            ) from exc

        self.torch = torch
        self.model = AutoModelForImageSegmentation.from_pretrained(
            settings.model_name, trust_remote_code=True
        )
        self.model.eval()
        self.model.float()  # checkpoint ships as fp16; force fp32 to match inputs
        self.device = self._resolve_device(settings.device)
        self.model.to(self.device)

        size = settings.input_size
        self.transform = transforms.Compose(
            [
                transforms.Resize((size, size)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )

    def _resolve_device(self, requested: str):
        if requested == "cuda" and not self.torch.cuda.is_available():
            raise RuntimeError(
                "CUDA was requested but is not available. Use --device cpu or install a CUDA-enabled PyTorch build."
            )
        return self.torch.device(requested)

    def mask(self, rgb: np.ndarray) -> np.ndarray:
        """Return a soft foreground mask, float32 0..1, at the input resolution."""
        image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
        tensor = self.transform(image).unsqueeze(0).to(self.device)
        with self.torch.no_grad():
            pred = self.model(tensor)[-1].sigmoid()
        pred = pred[0, 0].detach().float().cpu().numpy()
        if pred.shape != rgb.shape[:2]:
            pred_img = Image.fromarray(
                np.clip(pred * 255.0, 0, 255).astype(np.uint8), mode="L"
            ).resize((rgb.shape[1], rgb.shape[0]), Image.Resampling.BILINEAR)
            pred = np.asarray(pred_img, dtype=np.float32) / 255.0
        return np.clip(pred, 0.0, 1.0).astype(np.float32)