from __future__ import annotations

import base64
from io import BytesIO


VLM_SUPPORTED_IMAGE_MIME = {"image/jpeg", "image/png"}
JPEG_QUALITY = 90


def image_mime(data: bytes, content_type: str = "") -> str:
    if data.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if data.startswith(b"GIF87a") or data.startswith(b"GIF89a"):
        return "image/gif"
    if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "image/webp"
    raise RuntimeError("image_signature_invalid")


def make_image_data_uri(data: bytes, mime: str) -> str:
    return "data:" + mime + ";base64," + base64.b64encode(data).decode("ascii")


def normalize_image_for_vision(data: bytes, detected_mime: str) -> tuple[bytes, str, str]:
    if detected_mime == "image/jpeg":
        return data, "image/jpeg", "pass_through"
    if detected_mime == "image/png":
        return data, "image/png", "pass_through"
    if detected_mime != "image/webp":
        raise RuntimeError("image_mime_unsupported")
    try:
        from PIL import Image, ImageOps

        with Image.open(BytesIO(data)) as source:
            try:
                source.seek(0)
            except Exception:
                pass
            image = ImageOps.exif_transpose(source)
            has_alpha = image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info)
            if has_alpha:
                rgba = image.convert("RGBA")
                background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
                background.alpha_composite(rgba)
                image = background.convert("RGB")
            elif image.mode != "RGB":
                image = image.convert("RGB")
            output = BytesIO()
            image.save(output, format="JPEG", quality=JPEG_QUALITY, optimize=True)
            normalized = output.getvalue()
    except Exception as exc:
        raise RuntimeError("image_normalize_failed") from exc
    if not normalized.startswith(b"\xff\xd8\xff"):
        raise RuntimeError("image_normalize_invalid")
    return normalized, "image/jpeg", "webp_to_jpeg"


def normalize_image_bytes_to_data_uri(data: bytes, content_type: str = "") -> tuple[str, str, str, int, int]:
    detected_mime = image_mime(data, content_type)
    normalized, normalized_mime, action = normalize_image_for_vision(data, detected_mime)
    if normalized_mime not in VLM_SUPPORTED_IMAGE_MIME:
        raise RuntimeError("image_mime_unsupported")
    return make_image_data_uri(normalized, normalized_mime), detected_mime, normalized_mime, len(data), len(normalized)


def normalize_data_uri_for_vision(data_uri: str, max_bytes: int = 10 * 1024 * 1024) -> tuple[str, str, str, int, int]:
    value = str(data_uri or "").strip()
    prefix = "data:"
    marker = ";base64,"
    if not value.startswith(prefix) or marker not in value[:128]:
        raise RuntimeError("image_data_uri_invalid")
    header, encoded = value.split(marker, 1)
    if not header.lower().startswith("data:image/"):
        raise RuntimeError("image_data_uri_invalid")
    try:
        data = base64.b64decode(encoded, validate=True)
    except Exception as exc:
        raise RuntimeError("image_data_uri_decode_failed") from exc
    if not data or len(data) > max_bytes:
        raise RuntimeError("image_size_invalid")
    declared_mime = header[5:].split(";", 1)[0].strip().lower()
    return normalize_image_bytes_to_data_uri(data, declared_mime)
