"""Screenshot capture for macOS.

Reads a JPEG file written by the Swift host app (LockInBro.app), which holds
the Screen Recording TCC permission. The Swift app writes atomically to
/tmp/lockinbro_capture.jpg every ~5s using ScreenCaptureKit.

capture_screenshot() blocks up to 15 seconds waiting for a fresh file,
then raises RuntimeError if nothing arrives — so loop.py can handle it
gracefully without falling back to the screencapture CLI (which would fail
with a permissions error when called from a Python subprocess).
"""

from __future__ import annotations

import io
import os
import time

from PIL import Image

from argus.config import SCREENSHOT_JPEG_QUALITY, SCREENSHOT_MAX_WIDTH

_SWIFT_FRAME_PATH = "/tmp/lockinbro_capture.jpg"
_SWIFT_FRAME_MAX_AGE_S = 10.0   # treat as stale if older than this
_WAIT_TIMEOUT_S = 15.0          # how long to wait for Swift to write the first frame
_WAIT_POLL_S = 0.5              # polling interval while waiting


def _to_jpeg(img: Image.Image) -> bytes:
    """Convert a PIL Image to JPEG bytes, handling RGBA → RGB."""
    if img.width > SCREENSHOT_MAX_WIDTH:
        ratio = SCREENSHOT_MAX_WIDTH / img.width
        img = img.resize(
            (SCREENSHOT_MAX_WIDTH, int(img.height * ratio)),
            Image.LANCZOS,
        )
    if img.mode == "RGBA":
        img = img.convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=SCREENSHOT_JPEG_QUALITY)
    return buf.getvalue()


def _read_swift_frame() -> bytes | None:
    """Return JPEG bytes from the Swift-provided file if it exists and is fresh."""
    try:
        age = time.time() - os.path.getmtime(_SWIFT_FRAME_PATH)
        if age >= _SWIFT_FRAME_MAX_AGE_S:
            return None
        with open(_SWIFT_FRAME_PATH, "rb") as f:
            data = f.read()
        img = Image.open(io.BytesIO(data))
        img.load()
        return _to_jpeg(img)
    except Exception:
        return None


def capture_screenshot() -> bytes:
    """Return JPEG bytes for the current screen.

    Blocks up to _WAIT_TIMEOUT_S seconds waiting for the Swift app to write
    a fresh frame. Raises RuntimeError if no frame arrives in time.
    """
    deadline = time.time() + _WAIT_TIMEOUT_S
    while time.time() < deadline:
        frame = _read_swift_frame()
        if frame is not None:
            return frame
        time.sleep(_WAIT_POLL_S)

    raise RuntimeError(
        f"No screenshot received from LockInBro.app within {_WAIT_TIMEOUT_S}s. "
        "Make sure the app is running and has Screen Recording permission."
    )


def load_image_file(path: str) -> bytes:
    """Load an image file and return JPEG bytes (for testing)."""
    return _to_jpeg(Image.open(path))