"""Cap PNG dimensions to stay under provider per-image limits.

Claude rejects multi-image conversations where any image exceeds 2000px
on the long edge. Local vision models (LLaVA, Qwen-VL) work best at
≤1024px. **Default cap is 1600** — that's 400px of headroom under the
2000 API ceiling, deliberately wider than necessary. We tried 1920 first
and still saw intermittent rejections in production (the 80px margin was
swamped by rounding, accumulation across many images per turn, and one
stale-subprocess incident). 1600 is the new safe default; override with
`MCP_MAX_IMAGE_DIM` if you genuinely need higher resolution.

Three backends, tried in order — first one that works wins:

  1. **cv2** (OpenCV) — best quality, but requires the `[ar]` extra
     (~80 MB install). Available when the project is set up for
     AR/Vision work.
  2. **PIL** (Pillow) — lighter, often already pulled in by other
     deps. Quality slightly worse than cv2 on photo content; fine
     for UI screenshots.
  3. **sips** — macOS-native, zero Python deps. Always available on
     a developer Mac. The defaulting layer for the common case.

If all three fail (Linux box without cv2/PIL), `cap_image_in_place`
returns False AND logs a clear stderr warning so the user knows the
cap is non-functional. Together with the dispatcher's hard guard
(see `presentation/image_safety_net.py`), this guarantees the
2000px limit is never silently violated.

Originals stay on disk at `<path>.orig.png` so visual-diff workflows
have full resolution.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from importlib.util import find_spec
from pathlib import Path

from ..observability import warn

DEFAULT_MAX_DIM = 1600


def _max_dim() -> int:
    raw = os.environ.get("MCP_MAX_IMAGE_DIM", "")
    if not raw:
        return DEFAULT_MAX_DIM
    try:
        value = int(raw)
    except ValueError:
        return DEFAULT_MAX_DIM
    return value  # 0 disables capping


def prefer_original(path: Path) -> Path:
    """Return `<path>.orig.png` if it exists, else `path` unchanged.

    Visual-diff workflows (compare_screenshot, golden image checks)
    should always read from this so capping doesn't degrade their
    accuracy. Capping happens for vision-model context budget — diff
    math wants the full sensor.
    """
    original = path.with_suffix(".orig.png")
    return original if original.exists() else path


# ---- dimension probing (no deps) ----------------------------------------


def _read_png_dimensions(path: Path) -> tuple[int, int] | None:
    """Read PNG IHDR width/height without any image library. Returns None
    if the file isn't a valid PNG header."""
    try:
        with path.open("rb") as fh:
            header = fh.read(24)
    except OSError:
        return None
    if len(header) < 24 or header[:8] != b"\x89PNG\r\n\x1a\n":
        return None
    width = int.from_bytes(header[16:20], "big")
    height = int.from_bytes(header[20:24], "big")
    return (width, height)


# ---- per-backend resize implementations ---------------------------------


def _resize_cv2(path: Path, new_w: int, new_h: int) -> bool:
    if find_spec("cv2") is None:
        return False
    try:
        import cv2

        img = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if img is None:
            return False
        resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        # cv2's PNG writer defaults to compression level 1 (fast but big).
        # Force level 9 so post-cap files aren't ~3x larger than they
        # need to be. If PIL is also available, follow up with a palette
        # recompress for an additional 2-3x size drop on UI content.
        ok_write = bool(cv2.imwrite(
            str(path), resized, [cv2.IMWRITE_PNG_COMPRESSION, 9]
        ))
        if ok_write and find_spec("PIL") is not None:
            compress_png_in_place(path)  # palette pass; ignore failures
        return ok_write
    except Exception as exc:
        warn("image_cap_backend_failed", backend="cv2", error=str(exc))
        return False


def _resize_pil(path: Path, new_w: int, new_h: int) -> bool:
    if find_spec("PIL") is None:
        return False
    try:
        from PIL import Image

        with Image.open(path) as img:
            resized = img.resize((new_w, new_h), Image.LANCZOS)
            # Use _save_compressed_png so the post-resize file gets the
            # same compression treatment as the standalone compression
            # path applies. Big win for UI screenshots: 4-6x smaller
            # without visible loss.
            _save_compressed_png(resized, path)
        return True
    except Exception as exc:
        warn("image_cap_backend_failed", backend="PIL", error=str(exc))
        return False


def _save_compressed_png(img, path: Path) -> None:
    """Save `img` to `path` as a PNG using palette mode + max zlib compression.

    Why palette mode: UI screenshots have small color palettes (typically
    < 256 unique colors per frame). Converting truecolor → 256-color
    adaptive palette is visually lossless for UI content but slashes file
    size 3-5x. Combined with `optimize=True` + `compress_level=9` the
    typical post-cap UI screenshot drops from ~600 KB to ~150 KB.

    Falls back to plain RGB if palette conversion fails for any reason
    (e.g. transparency edge cases we haven't tested).

    Why this matters: even with the dimension cap at 1600 px, agents
    accumulate dozens of screenshots in a long e2e flow. The upstream
    API enforces a 32 MB total request size — 50 truecolor PNGs
    @ 600 KB each = 30 MB, over the limit. Palette mode keeps the same
    50 shots well under 10 MB.
    """
    from PIL import Image

    try:
        # `RGBA` images sometimes survive better through quantize than
        # convert("P"). The `Image.Quantize` API picks an adaptive palette
        # from the actual image pixels.
        if img.mode in ("RGB", "RGBA"):
            palette_img = img.convert("RGB").quantize(
                colors=256, method=Image.Quantize.MEDIANCUT
            )
        else:
            palette_img = img.convert("P", palette=Image.Palette.ADAPTIVE, colors=256)
        palette_img.save(
            path, format="PNG", optimize=True, compress_level=9
        )
    except Exception:
        # Last-ditch: write as-is with max compression.
        img.save(path, format="PNG", optimize=True, compress_level=9)


def compress_png_in_place(path: Path) -> bool:
    """Re-encode an existing PNG with palette + max compression.

    Use this on PNGs that are already at acceptable dimensions but too
    large in bytes (e.g. golden images, or images produced by tools that
    don't route through `cap_image_in_place`). Returns True if the file
    was rewritten (it's smaller); False if PIL isn't available or the
    re-encode failed.
    """
    if find_spec("PIL") is None:
        return False
    try:
        from PIL import Image

        original_size = path.stat().st_size
        with Image.open(path) as img:
            img.load()  # decode before we close the source handle
            _save_compressed_png(img, path)
        new_size = path.stat().st_size
        return new_size < original_size
    except Exception as exc:
        warn("png_compression_failed", path=str(path), error=str(exc))
        return False


def _resize_sips(path: Path, cap: int) -> bool:
    """macOS `sips -Z` resizes the longest dimension to `cap` in place.

    Doesn't need explicit new_w/new_h — sips does the math itself. Only
    available on macOS; returns False on other platforms.
    """
    if not shutil.which("sips"):
        return False
    try:
        result = subprocess.run(
            ["sips", "-Z", str(cap), str(path)],
            capture_output=True,
            timeout=10,
        )
        return result.returncode == 0
    except (subprocess.TimeoutExpired, OSError) as exc:
        warn("image_cap_backend_failed", backend="sips", error=str(exc))
        return False


# ---- the public entry point ---------------------------------------------


def cap_image_in_place(path: Path, max_dim: int | None = None) -> bool:
    """Resize `path` so its long edge ≤ max_dim. Returns True if resized.

    Idempotent on already-capped images (no-op when both dims are
    already within bounds). Preserves the original at `<path>.orig.png`
    on the first call.

    Tries cv2 → PIL → sips. If all three fail (and the image is over
    cap), logs to stderr — the dispatcher's hard guard will refuse to
    return the path in that case.
    """
    cap = _max_dim() if max_dim is None else max_dim
    if cap <= 0:
        return False
    dims = _read_png_dimensions(path)
    if dims is None:
        return False
    w, h = dims
    longest = max(w, h)
    if longest <= cap:
        return False

    # Snapshot the original on first cap.
    original = path.with_suffix(".orig.png")
    if not original.exists():
        try:
            shutil.copy2(path, original)
        except OSError as exc:
            warn("image_cap_original_save_failed", path=str(path), error=str(exc))
            # Continue anyway — losing the original is worse than failing
            # to cap, but capping is what protects the conversation.

    scale = cap / longest
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))

    # Try backends in order. First successful one wins.
    for backend in (
        lambda: _resize_cv2(path, new_w, new_h),
        lambda: _resize_pil(path, new_w, new_h),
        lambda: _resize_sips(path, cap),
    ):
        if backend():
            return True

    # All three failed.
    warn(
        "image_cap_all_backends_failed",
        path=str(path),
        width=w,
        height=h,
        cap=cap,
        fix="install cv2 (uv pip install -e '.[ar]') or Pillow, or run on macOS",
    )
    return False


def available_backends() -> tuple[str, ...]:
    """Which cap backends are usable in this process, in priority order.

    Surfaced by `mcp_ping` and the boot self-check log so a stale
    subprocess (or a Linux box missing cv2/PIL) is visible in one call
    instead of failing silently. If this returns `()`, capping is
    non-functional and the safety net will refuse oversized PNGs.
    """
    found: list[str] = []
    if find_spec("cv2") is not None:
        found.append("cv2")
    if find_spec("PIL") is not None:
        found.append("PIL")
    if shutil.which("sips"):
        found.append("sips")
    return tuple(found)


def is_within_cap(path: Path, max_dim: int | None = None) -> bool:
    """Verification helper: True if `path` is a PNG whose long edge ≤ cap.

    Used by the dispatcher seatbelt to verify cap success before
    returning a path to the agent. Returns True for non-PNG files
    (they don't trigger the multi-image limit).
    """
    cap = _max_dim() if max_dim is None else max_dim
    if cap <= 0:
        return True
    dims = _read_png_dimensions(path)
    if dims is None:
        # Not a PNG or unreadable — let it pass; multi-image limit only
        # affects images Claude can decode.
        return True
    return max(dims) <= cap
