"""Background removal service using rembg + BiRefNet."""

import io
import logging

from PIL import Image
from rembg import new_session, remove

from app.config import settings

logger = logging.getLogger(__name__)

MAX_INPUT_SIZE = 2048


class SegmentationService:
    """Stateless background removal using rembg with BiRefNet model."""

    def __init__(self):
        import onnxruntime as ort

        # onnxruntime-gpu wheel ships without CUDA libs; load them from
        # the nvidia-*-cu12 pip packages before session creation.
        if hasattr(ort, "preload_dlls"):
            try:
                ort.preload_dlls()
            except Exception as e:
                logger.warning("ort.preload_dlls() failed: %s", e)

        model = settings.SEGMENTATION_MODEL
        available = ort.get_available_providers()
        logger.info(
            "onnxruntime device=%s available=%s", ort.get_device(), available
        )

        requested = [
            p.strip()
            for p in settings.SEGMENTATION_PROVIDERS.split(",")
            if p.strip()
        ] or None

        logger.info(
            "Loading rembg session (model=%s, providers=%s)...",
            model,
            requested or "auto",
        )
        if requested is not None:
            self._session = new_session(model_name=model, providers=requested)
        else:
            self._session = new_session(model_name=model)

        active = self._session.inner_session.get_providers()
        logger.info("rembg session active providers=%s", active)
        if requested and active[: len(requested)] != requested:
            logger.warning(
                "Requested providers %s but session uses %s — CUDA libs may have failed to load",
                requested,
                active,
            )
        elif "CUDAExecutionProvider" in available and "CUDAExecutionProvider" not in active:
            logger.warning(
                "CUDAExecutionProvider is available but session fell back to %s",
                active,
            )

    def remove_background(
        self, image_bytes: bytes, center_crop: bool = False
    ) -> bytes:
        """Remove background from image bytes.

        Args:
            image_bytes: Raw image bytes (JPEG, PNG, WebP, etc.)
            center_crop: If True, crop to 9:16 with person centered and
                         minimal top/bottom margins.

        Returns:
            PNG bytes with transparent background.
        """
        resized = self._resize_for_segmentation(image_bytes)
        result: bytes = remove(
            data=resized,
            session=self._session,
            post_process_mask=True,
        )
        if center_crop:
            img = Image.open(io.BytesIO(result)).convert("RGBA")
            img = self._center_crop_person(img)
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            result = buf.getvalue()
        return result

    @staticmethod
    def _center_crop_person(
        img: Image.Image,
        pad_top_ratio: float = 0.05,
        pad_bottom_ratio: float = 0.00,
        pad_horizontal_ratio: float = 0.02,
    ) -> Image.Image:
        """Crop RGBA image to 9:16 with person centered, minimal margins.

        Detects the person region from the alpha channel bounding box,
        adds asymmetric vertical padding (top 5%, bottom 10%) and
        horizontal padding (5%), then fits to 9:16 aspect ratio.
        The person is never cropped — if the person is wider than 9:16
        allows, the frame expands vertically to maintain the ratio.
        Areas outside the original image are filled with transparency.
        """
        alpha = img.getchannel("A")
        bbox = alpha.getbbox()
        if bbox is None:
            logger.warning("center_crop: no non-transparent pixels found, returning as-is")
            return img

        left, top, right, bottom = bbox
        person_w = right - left
        person_h = bottom - top

        # Padded person dimensions (person must fully fit)
        padded_w = person_w + int(person_w * pad_horizontal_ratio) * 2
        padded_h = person_h + int(person_h * pad_top_ratio) + int(person_h * pad_bottom_ratio)

        # Determine final size: fit 9:16 while containing the padded person
        # Try height-driven: width = padded_h * 9/16
        width_from_h = int(padded_h * 9 / 16)
        if width_from_h >= padded_w:
            # Height-driven fits — person width fits within 9:16 frame
            target_w = width_from_h
            target_h = padded_h
        else:
            # Width-driven — person is too wide, expand height to keep 9:16
            target_w = padded_w
            target_h = int(padded_w * 16 / 9)

        # Person center
        person_center_x = (left + right) // 2
        person_center_y = (top + bottom) // 2

        # Canvas region in original image coordinates
        # Vertical: center person in canvas, then shift up slightly
        # to account for asymmetric padding (more space on top than bottom)
        pad_bias = int(person_h * (pad_top_ratio - pad_bottom_ratio) / 2)
        canvas_top = person_center_y - target_h // 2 - pad_bias
        canvas_bottom = canvas_top + target_h

        # Horizontal: center on person
        canvas_left = person_center_x - target_w // 2
        canvas_right = canvas_left + target_w

        # Create transparent canvas
        canvas = Image.new("RGBA", (target_w, target_h), (0, 0, 0, 0))

        # Compute overlap between canvas region and original image
        src_left = max(0, canvas_left)
        src_right = min(img.width, canvas_right)
        src_top = max(0, canvas_top)
        src_bottom = min(img.height, canvas_bottom)

        if src_left >= src_right or src_top >= src_bottom:
            return canvas

        region = img.crop((src_left, src_top, src_right, src_bottom))
        dst_x = src_left - canvas_left
        dst_y = src_top - canvas_top
        canvas.paste(region, (dst_x, dst_y))
        return canvas

    @staticmethod
    def _resize_for_segmentation(
        image_bytes: bytes,
        max_size: int = MAX_INPUT_SIZE,
    ) -> bytes:
        """Resize image if too large, preserving format as PNG for alpha support."""
        img = Image.open(io.BytesIO(image_bytes))
        w, h = img.size
        if max(w, h) <= max_size:
            return image_bytes

        ratio = max_size / max(w, h)
        new_size = (int(w * ratio), int(h * ratio))
        img = img.resize(new_size, Image.LANCZOS)

        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()