"""Validated single-image input for Decisions. Images stay out of the text state."""
import base64
import binascii
import io

from PIL import Image, UnidentifiedImageError

from decisions_api import ApiError

IMG_MARK = "<<IMG>>"
IMG_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
MAX_IMAGE_BYTES = 4 * 1024 * 1024
MAX_IMAGE_PIXELS = 12_000_000


def image_input(body):
    """Accept base64 or a data URI; never read server paths or fetch remote URLs."""
    value = body.get("image_data")
    if value is None:
        return None
    if not isinstance(value, str) or not value:
        raise ApiError(400, "image_data must be one base64 image or image data URI")
    if value.startswith("data:"):
        header, sep, value = value.partition(",")
        if not sep or not header.startswith("data:image/") or not header.endswith(";base64"):
            raise ApiError(400, "image_data must use an image base64 data URI")
    if len(value) > 4 * ((MAX_IMAGE_BYTES + 2) // 3):
        raise ApiError(413, "image_data exceeds 4 MiB decoded")
    try:
        raw = base64.b64decode(value, validate=True)
        if len(raw) > MAX_IMAGE_BYTES:
            raise ApiError(413, "image_data exceeds 4 MiB decoded")
        with Image.open(io.BytesIO(raw)) as im:
            if im.format not in {"JPEG", "PNG", "WEBP"} or getattr(im, "n_frames", 1) != 1:
                raise ApiError(400, "image_data must be a single JPEG, PNG or WebP image")
            if im.width * im.height > MAX_IMAGE_PIXELS:
                raise ApiError(413, "image_data exceeds 12 million pixels")
            im.verify()
    except (binascii.Error, ValueError, OSError, UnidentifiedImageError, Image.DecompressionBombError) as e:
        raise ApiError(400, "image_data is not a valid base64 image") from e
    return value


def image_premise(premise):
    if premise.count(IMG_MARK) > 1 or "<|image_pad|>" in premise or "<|vision_start|>" in premise:
        raise ApiError(400, "Use at most one <<IMG>> marker in the image premise")
    return premise.replace(IMG_MARK, IMG_BLOCK) if IMG_MARK in premise else premise.rstrip() + " " + IMG_BLOCK