from __future__ import annotations

import base64
import binascii
from typing import Any

_SUPPORTED_INLINE_IMAGE_MIME_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp"}


def extract_inline_image_candidates(artifact: dict[str, Any]) -> list[dict[str, Any]]:
    """Extract DATA_INLINE parts that look like image payloads, even before support checks."""

    image_parts: list[dict[str, Any]] = []
    for part in artifact.get("parts") or []:
        if not isinstance(part, dict):
            continue
        if part.get("type") != "DATA_INLINE":
            continue
        mime_type = part.get("mimeType")
        if not isinstance(mime_type, str) or not mime_type.strip().lower().startswith("image/"):
            continue
        image_parts.append(part)
    return image_parts


def filter_sendable_inline_image_parts(
    image_parts: list[dict[str, Any]],
) -> tuple[list[dict[str, str]], list[dict[str, Any]]]:
    """Separate inline image parts into provider-sendable payloads and skipped diagnostics."""

    sendable_parts: list[dict[str, str]] = []
    skipped_parts: list[dict[str, Any]] = []

    for index, part in enumerate(image_parts):
        mime_type = part["mimeType"].strip().lower()
        if mime_type not in _SUPPORTED_INLINE_IMAGE_MIME_TYPES:
            # Future improvement: rasterize unsupported image formats such as SVG to PNG before attaching them.
            skipped_parts.append({"index": index, "reason": "unsupported_image_mime_type_for_llm"})
            continue
        data_base64 = part.get("dataBase64")
        if not isinstance(data_base64, str) or not data_base64.strip():
            skipped_parts.append({"index": index, "reason": "invalid_image_data"})
            continue
        if not is_valid_image_base64(data_base64):
            skipped_parts.append({"index": index, "reason": "invalid_image_data"})
            continue
        # Emit the normalized MIME value so the provider sees the same shape we validated locally.
        sendable_parts.append(
            {
                "dataBase64": data_base64,
                "mimeType": mime_type,
            }
        )

    return sendable_parts, skipped_parts


def is_valid_image_base64(data_base64: str) -> bool:
    """Return whether the payload is decodable base64 after whitespace normalization."""

    # Remove surrounding and embedded ASCII whitespace before validating the base64 payload.
    normalized_data = "".join(data_base64.strip().split())
    try:
        base64.b64decode(normalized_data, validate=True)
    except (ValueError, binascii.Error):
        return False
    return True


def image_read_status(
    *,
    image_parts: list[dict[str, str]],
    sendable_parts: list[dict[str, str]],
    skipped_parts: list[dict[str, Any]],
) -> str:
    """Summarize the overall image-read outcome for the textual tool result."""

    if not image_parts:
        return "ok"
    if len(sendable_parts) == len(image_parts) and not skipped_parts:
        return "ok"
    if sendable_parts and skipped_parts:
        return "partial"
    if skipped_parts and all(item["reason"] == "unsupported_image_mime_type_for_llm" for item in skipped_parts):
        return "unsupported"
    if not sendable_parts:
        return "invalid"
    return "partial"
