"""Turn a raw product photo into a centered cutout on a square canvas."""

from __future__ import annotations

import logging
from io import BytesIO
from pathlib import PurePosixPath

from PIL import Image, ImageOps

logger = logging.getLogger(__name__)

CANVAS_SIZE = 1200
CUTOUT_MAX_SIDE = 1600
PADDING = 0.1
ALPHA_THRESHOLD = 128

_session = None


def _rembg_session():
    global _session
    if _session is None:
        from rembg import new_session

        _session = new_session("u2net")
    return _session


def open_image(data: bytes) -> Image.Image:
    image = Image.open(BytesIO(data))
    image.load()
    image = ImageOps.exif_transpose(image) or image
    if getattr(image, "n_frames", 1) > 1:
        image.seek(0)
        image = image.copy()
    return image


def downscale(image: Image.Image, max_side: int) -> Image.Image:
    width, height = image.size
    longest = max(width, height)
    if longest <= max_side:
        return image
    scale = max_side / longest
    return image.resize(
        (max(1, int(width * scale)), max(1, int(height * scale))),
        Image.Resampling.LANCZOS,
    )


def harden_cutout(image: Image.Image, threshold: int = ALPHA_THRESHOLD) -> Image.Image:
    """Drop faint shadow/halo pixels rembg leaves around the subject."""
    rgba = image.convert("RGBA")
    red, green, blue, alpha = rgba.split()
    alpha = alpha.point(lambda value: value if value >= threshold else 0)
    return Image.merge("RGBA", (red, green, blue, alpha))


def cutout_subject(image: Image.Image) -> Image.Image:
    """Knock out the background. Falls back to the original on failure."""
    rgba = downscale(image.convert("RGBA"), CUTOUT_MAX_SIDE)
    try:
        from rembg import remove

        buf = BytesIO()
        rgba.save(buf, format="PNG")
        result = remove(
            buf.getvalue(),
            session=_rembg_session(),
            post_process_mask=True,
        )
        return harden_cutout(Image.open(BytesIO(result)).convert("RGBA"))
    except Exception:
        logger.exception("product photo cutout failed; using original")
        return rgba


def center_on_square(
    image: Image.Image,
    *,
    size: int = CANVAS_SIZE,
    padding: float = PADDING,
) -> Image.Image:
    rgba = image.convert("RGBA")
    alpha = rgba.getchannel("A")
    mask = alpha.point(lambda value: 255 if value > ALPHA_THRESHOLD else 0)
    bbox = mask.getbbox()
    canvas = Image.new("RGBA", (size, size), (0, 0, 0, 0))
    if bbox is None:
        return canvas
    cropped = rgba.crop(bbox)
    max_inner = max(1, int(size * (1 - 2 * padding)))
    width, height = cropped.size
    scale = min(max_inner / width, max_inner / height)
    new_w = max(1, int(width * scale))
    new_h = max(1, int(height * scale))
    resized = cropped.resize((new_w, new_h), Image.Resampling.LANCZOS)
    x = (size - new_w) // 2
    y = (size - new_h) // 2
    canvas.paste(resized, (x, y), resized)
    return canvas


def encode_png(image: Image.Image) -> bytes:
    buf = BytesIO()
    image.save(buf, format="PNG", optimize=True)
    return buf.getvalue()


def prepare_product_photo(data: bytes) -> tuple[bytes, str]:
    """Return a PNG cutout on a square canvas, plus content type."""
    image = open_image(data)
    cut = cutout_subject(image)
    framed = center_on_square(cut)
    return encode_png(framed), "image/png"


def product_image_filename(original: str) -> str:
    stem = PurePosixPath(original or "product").stem.strip() or "product"
    return f"{stem[:200]}.png"