"""Image format detection, normalization and region cropping for vision tools.

Everything here runs BEFORE an image is base64-embedded. A vision tool result
is baked into immutable conversation history and re-sent every turn, so an
unsupported media type or corrupt bytes would wedge the session with a
non-retryable 400 on every resume — normalization must happen up front.
"""

from __future__ import annotations

import logging
import uuid
from io import BytesIO
from pathlib import Path
from typing import Any, Optional

from hermes_constants import get_hermes_dir

logger = logging.getLogger("tools.vision_tools")

_EXTENSION_MIME_TYPES = {
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".webp": "image/webp",
    ".svg": "image/svg+xml",
}

# Media types the major vision providers (Anthropic in particular) accept
# inline. SVG/BMP/TIFF are rejected with a non-retryable 400.
_ANTHROPIC_SUPPORTED_MEDIA_TYPES = frozenset(
    {"image/jpeg", "image/png", "image/gif", "image/webp"}
)


def _determine_mime_type(image_path: Path) -> str:
    """MIME type from file extension (defaults to image/jpeg)."""
    return _EXTENSION_MIME_TYPES.get(image_path.suffix.lower(), "image/jpeg")


def _detect_image_mime_type_from_bytes(data: bytes) -> Optional[str]:
    """Magic-byte MIME sniff (authoritative; no extension trust).

    Returns ``None`` for anything without a recognized header — including SVG,
    which has no magic bytes (the resolver sniffs ``<svg`` and passes it
    through for rasterization).
    """
    header = data[:64]
    if header.startswith(b"\x89PNG\r\n\x1a\n"):
        # Magic bytes alone are insufficient: reject corrupt PNGs before they
        # can be embedded. Pillow is optional — without it fall back to
        # header-only sniffing; only an actual failed verify() rejects.
        try:
            from PIL import Image
        except ImportError:
            return "image/png"
        try:
            with Image.open(BytesIO(data)) as image:
                image.verify()
        except Exception:
            return None
        return "image/png"
    if header.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if header.startswith((b"GIF87a", b"GIF89a")):
        return "image/gif"
    if header.startswith(b"BM"):
        return "image/bmp"
    if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
        return "image/webp"
    return None


def _supported_media_types() -> frozenset:
    """Formats the ACTIVE main model's server can decode.

    The managed llama-server decodes with stb_image — no WebP — and an
    undecodable image part fails SILENTLY (the model confabulates), so the set
    is narrowed there and normalization converts those formats to PNG.
    """
    try:
        from agent.auxiliary_client import _runtime_main_value
        from hermes_cli.local_runtime.capabilities import (
            ACCEPTED_IMAGE_MIMES,
            is_managed_provider,
        )

        if is_managed_provider(
                str(_runtime_main_value("provider") or ""),
                str(_runtime_main_value("base_url") or "")):
            return ACCEPTED_IMAGE_MIMES
    except Exception:  # noqa: BLE001 — best-effort narrowing only
        pass
    return _ANTHROPIC_SUPPORTED_MEDIA_TYPES


def _rasterize_svg_to_png(svg_path: Path, out_path: Path) -> bool:
    """Best-effort SVG → PNG via cairosvg, svglib+reportlab, rsvg-convert, inkscape (all soft deps)."""
    def _ok() -> bool:
        return out_path.exists() and out_path.stat().st_size > 0

    try:
        import cairosvg  # type: ignore
        cairosvg.svg2png(url=str(svg_path), write_to=str(out_path))
        return _ok()
    except Exception:
        pass
    try:
        from svglib.svglib import svg2rlg  # type: ignore
        from reportlab.graphics import renderPM  # type: ignore
        drawing = svg2rlg(str(svg_path))
        if drawing is not None:
            renderPM.drawToFile(drawing, str(out_path), fmt="PNG")
            return _ok()
    except Exception:
        pass
    import shutil
    import subprocess
    for cmd in (
        ["rsvg-convert", "-o", str(out_path), str(svg_path)],
        ["inkscape", str(svg_path), "--export-type=png",
         f"--export-filename={out_path}"],
    ):
        if shutil.which(cmd[0]):
            try:
                subprocess.run(
                    cmd, check=True, capture_output=True, timeout=30,
                    stdin=subprocess.DEVNULL,
                )
                if _ok():
                    return True
            except Exception:
                continue
    return False


def _normalize_to_supported_image(
    image_path: Path, detected_mime: str
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
    """Ensure an image is in a provider-supported format.

    Returns ``(path, mime, error)``: the input unchanged when already
    supported; ``(new_png_path, "image/png", None)`` after conversion — a temp
    file the CALLER must clean up; ``(None, None, message)`` when impossible.
    SVG is rasterized; other Pillow-readable rasters (BMP, TIFF) re-encode to PNG.
    """
    if detected_mime in _supported_media_types():
        return image_path, detected_mime, None

    out_dir = get_hermes_dir("cache/vision", "temp_vision_images")
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"converted_{uuid.uuid4()}.png"

    if detected_mime == "image/svg+xml":
        if _rasterize_svg_to_png(image_path, out_path):
            return out_path, "image/png", None
        return (
            None,
            None,
            "This is an SVG, which vision models cannot read directly, and no "
            "SVG rasterizer is installed (tried cairosvg, svglib, rsvg-convert, "
            "inkscape). Convert the SVG to PNG first — e.g. open it in a browser "
            "and screenshot it, or install a rasterizer "
            "(`pip install cairosvg`) — then re-run vision_analyze on the PNG.",
        )

    try:
        from PIL import Image as _PILImage
        with _PILImage.open(image_path) as _img:
            if _img.mode not in ("RGB", "RGBA", "L"):
                _img = _img.convert("RGBA")
            _img.save(out_path, format="PNG")
        if out_path.exists() and out_path.stat().st_size > 0:
            return out_path, "image/png", None
    except Exception as _exc:
        logger.warning("Failed to normalize %s image to PNG: %s",
                       detected_mime, _exc)
    return (
        None,
        None,
        f"Image format {detected_mime!r} is not supported by the vision API "
        f"and could not be converted to PNG (install Pillow for raster "
        f"conversion). Convert it to PNG or JPEG and try again.",
    )


# Full raster validation runs on untrusted images in a shared CPU executor:
# bound animated-image work by frame count AND total decoded area so a compact
# file cannot monopolize a worker with unbounded frames.
_VISION_MAX_VALIDATED_FRAME_COUNT = 100
_VISION_MAX_VALIDATED_AGGREGATE_PIXELS = 100_000_000


def _validate_raster_image_decodable(
    image_path: Path,
    max_frames: int = _VISION_MAX_VALIDATED_FRAME_COUNT,
    max_pixels: int = _VISION_MAX_VALIDATED_AGGREGATE_PIXELS,
) -> Optional[str]:
    """Return an error unless Pillow can fully decode every frame.

    Header sniffing and ``Image.open`` only inspect containers: a timed-out
    download can look like a valid PNG with a truncated pixel stream. Without
    Pillow the image passes unvalidated rather than rejecting everything.
    """
    try:
        from PIL import Image as _PILImage
        from PIL import ImageSequence as _PILImageSequence
    except ImportError:
        return None
    try:
        with _PILImage.open(image_path) as image:
            image.verify()
        with _PILImage.open(image_path) as image:
            validated_pixels = 0
            for frame_number, frame in enumerate(
                _PILImageSequence.Iterator(image), start=1
            ):
                if frame_number > max_frames:
                    return (
                        "Image validation rejected animation: "
                        f"frame {frame_number} exceeds the maximum "
                        f"{max_frames} validated frames."
                    )
                next_validated_pixels = validated_pixels + frame.width * frame.height
                if next_validated_pixels > max_pixels:
                    return (
                        "Image validation rejected animation: aggregate decoded "
                        f"pixel count would reach {next_validated_pixels} at frame "
                        f"{frame_number}, exceeding the maximum "
                        f"{max_pixels}."
                    )
                frame.load()
                validated_pixels = next_validated_pixels
    except Exception as exc:
        return f"Image could not be fully decoded: {exc}"
    return None


def _image_exceeds_dimension(image_path: Path, max_dimension: int) -> bool:
    """True if the longest side exceeds ``max_dimension`` px.

    Anthropic enforces an 8000px per-side cap independently of the byte cap.
    Returns False (no forced resize) without Pillow or on unreadable files —
    a missing soft dependency must never break the embed path.
    """
    try:
        from PIL import Image as _PILImage
        with _PILImage.open(image_path) as _img:
            return max(_img.size) > max_dimension
    except Exception:
        return False


def _crop_image_region(
    image_path: Path,
    region: Any,
    offset_out: Optional[dict] = None,
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
    """Crop to ``region`` = [x1, y1, x2, y2] (original-image pixels).

    Applied BEFORE downscaling so the crop gets the full resolution budget.
    Coordinates clamp to the image bounds; a zero-area/inverted region is
    rejected with an error naming the real dimensions. Returns
    ``(cropped_temp_path, mime, None)`` — caller owns cleanup — or
    ``(None, None, error)``. Ported from QwenLM/qwen-code zoom-image.ts (Apache-2.0).
    """
    try:
        from PIL import Image
    except ImportError:
        return None, None, (
            "region cropping requires Pillow (`pip install Pillow`); "
            "retry without the region parameter."
        )

    if (
        not isinstance(region, (list, tuple))
        or len(region) != 4
        or not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in region)
    ):
        return None, None, (
            "Invalid region: expected [x1, y1, x2, y2] as four numbers "
            "(pixel coordinates in the original image)."
        )

    try:
        with Image.open(image_path) as img:
            width, height = img.size
            x1, y1, x2, y2 = (int(v) for v in region)
            cx1 = max(0, min(x1, width))
            cy1 = max(0, min(y1, height))
            cx2 = max(0, min(x2, width))
            cy2 = max(0, min(y2, height))
            if cx2 <= cx1 or cy2 <= cy1:
                return None, None, (
                    f"Invalid region [{x1}, {y1}, {x2}, {y2}]: crops to zero "
                    f"area after clamping to the image bounds. The image is "
                    f"{width}x{height} px — pick x1<x2 and y1<y2 inside "
                    f"[0, 0, {width}, {height}]."
                )
            cropped = img.crop((cx1, cy1, cx2, cy2))
            if offset_out is not None:
                offset_out.update(x=cx1, y=cy1, width=cx2 - cx1, height=cy2 - cy1)
            out_path = image_path.with_name(
                f"{image_path.stem}_region_{uuid.uuid4().hex[:8]}.png"
            )
            if cropped.mode not in ("RGB", "RGBA", "L", "LA", "P"):
                cropped = cropped.convert("RGB")
            cropped.save(out_path, format="PNG")
            return out_path, "image/png", None
    except Exception as exc:
        return None, None, f"Failed to crop region: {exc}"