"""Decode and re-encode small raster uploads before they leave MetalCircle."""
from io import BytesIO
import warnings

from PIL import Image, ImageOps


MAX_SCREENSHOT_BYTES = 4 * 1024 * 1024
MAX_SCREENSHOT_PIXELS = 16_000_000
MAX_SCREENSHOT_OUTPUT_BYTES = 8 * 1024 * 1024


class ScreenshotError(ValueError):
    """Stable validation code; never includes the uploaded filename or bytes."""
    def __init__(self, code):
        self.code = code
        super().__init__(code)


def sanitize_screenshot(contents):
    if not contents:
        raise ScreenshotError('empty')
    if len(contents) > MAX_SCREENSHOT_BYTES:
        raise ScreenshotError('too_large')

    try:
        with warnings.catch_warnings():
            warnings.simplefilter('error', Image.DecompressionBombWarning)
            with Image.open(BytesIO(contents)) as candidate:
                if candidate.format not in {'JPEG', 'PNG', 'WEBP'}:
                    raise ScreenshotError('unsupported')
                if candidate.width * candidate.height > MAX_SCREENSHOT_PIXELS:
                    raise ScreenshotError('too_large')
                if getattr(candidate, 'n_frames', 1) != 1:
                    raise ScreenshotError('unsupported')
                candidate.verify()
            with Image.open(BytesIO(contents)) as candidate:
                if candidate.width * candidate.height > MAX_SCREENSHOT_PIXELS:
                    raise ScreenshotError('too_large')
                image = ImageOps.exif_transpose(candidate)
                image.load()
                has_alpha = image.mode in {'RGBA', 'LA'} or (
                    image.mode == 'P' and 'transparency' in image.info
                )
                if image.mode not in {'RGB', 'RGBA'}:
                    image = image.convert('RGBA' if has_alpha else 'RGB')
                image.thumbnail((1600, 1600))
                output = BytesIO()
                if image.mode == 'RGBA':
                    image.save(output, format='PNG', optimize=True)
                    extension, media_type = 'png', 'image/png'
                else:
                    image.save(output, format='JPEG', quality=85, optimize=True)
                    extension, media_type = 'jpg', 'image/jpeg'
                safe_contents = output.getvalue()
    except ScreenshotError:
        raise
    except Exception:
        # Decoders may raise different exception types for malformed inputs.
        # Fail closed without exposing decoder details to the caller.
        raise ScreenshotError('invalid') from None

    if len(safe_contents) > MAX_SCREENSHOT_OUTPUT_BYTES:
        raise ScreenshotError('too_large')
    return safe_contents, extension, media_type