"""Screen capture tools."""

import asyncio
import io
import json
import time

import numpy as np
from fastmcp import Context
from fastmcp.exceptions import ToolError
from fastmcp.tools.tool import ToolResult
from fastmcp.utilities.types import Image
from PIL import Image as PILImage

from mcvnc._state import (
    SCREENSHOT_MAX_WIDTH,
    SCREENSHOT_TIMEOUT,
    WAIT_CHANGE_TIMEOUT,
    VncSession,
    get_state,
    mcp,
)


def _rgba_to_png(pixels, max_width: int = 0) -> bytes:
    """Convert a numpy RGBA array to PNG bytes, with optional downscale."""
    img = PILImage.fromarray(pixels)
    if max_width > 0 and img.width > max_width:
        ratio = max_width / img.width
        new_height = int(img.height * ratio)
        img = img.resize((max_width, new_height), PILImage.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return buf.getvalue()


@mcp.tool()
async def vnc_screenshot(
    ctx: Context,
    session_id: str,
    max_width: int = SCREENSHOT_MAX_WIDTH,
) -> ToolResult:
    """Capture a screenshot of the remote desktop.

    Returns metadata JSON and an inline PNG image. Use this to see what's
    on the remote screen before interacting.

    Args:
        session_id: VNC session to capture.
        max_width: Downscale to this width (0 = original size).
    """
    state = get_state(ctx)
    session = state.get(session_id)

    async with session.lock:
        try:
            pixels = await asyncio.wait_for(
                session.client.screenshot(), timeout=SCREENSHOT_TIMEOUT
            )
        except asyncio.TimeoutError:
            raise ToolError(
                f"Screenshot timed out after {SCREENSHOT_TIMEOUT}s. "
                f"The VNC server may be unresponsive. Try again or disconnect."
            )
        except Exception as e:
            raise ToolError(f"Screenshot failed: {e}")

        # Update cached dimensions inside the lock
        session.width = session.client.video.width
        session.height = session.client.video.height

    png_bytes = _rgba_to_png(pixels, max_width)
    meta = {
        "session_id": session_id,
        "width": session.width,
        "height": session.height,
        "png_size_bytes": len(png_bytes),
    }
    if max_width > 0 and session.width > max_width:
        ratio = max_width / session.width
        meta["scaled_width"] = max_width
        meta["scaled_height"] = int(session.height * ratio)
        meta["note"] = "Coordinates in other tools use original resolution, not scaled."

    return ToolResult(content=[json.dumps(meta), Image(data=png_bytes, format="png")])


@mcp.tool()
async def vnc_screen_info(ctx: Context, session_id: str) -> str:
    """Get screen dimensions without capturing a screenshot.

    Args:
        session_id: VNC session to query.
    """
    state = get_state(ctx)
    session = state.get(session_id)
    return json.dumps({
        "session_id": session_id,
        "width": session.width,
        "height": session.height,
        "desktop_name": session.client.video.name,
    })


def _validate_region(
    session: VncSession, x: int, y: int, width: int, height: int
) -> None:
    """Validate region bounds against the desktop dimensions."""
    if x < 0 or y < 0:
        raise ToolError(
            f"Region x and y must be non-negative (got x={x}, y={y})"
        )
    if width <= 0 or height <= 0:
        raise ToolError(
            f"Region width and height must be positive (got {width}x{height})"
        )
    if x + width > session.width or y + height > session.height:
        raise ToolError(
            f"Region {x},{y} {width}x{height} exceeds desktop bounds "
            f"({session.width}x{session.height})"
        )


@mcp.tool()
async def vnc_screenshot_region(
    ctx: Context,
    session_id: str,
    x: int,
    y: int,
    width: int,
    height: int,
    max_width: int = SCREENSHOT_MAX_WIDTH,
) -> ToolResult:
    """Capture a region of the remote desktop instead of the full screen.

    Use this when you only need to see a dialog, menu, or specific area.
    More efficient than a full screenshot when you know where to look.

    Args:
        session_id: VNC session to capture.
        x: Left edge of the region (pixels from left).
        y: Top edge of the region (pixels from top).
        width: Width of the region in pixels.
        height: Height of the region in pixels.
        max_width: Downscale to this width (0 = original size).
    """
    state = get_state(ctx)
    session = state.get(session_id)
    _validate_region(session, x, y, width, height)

    async with session.lock:
        try:
            pixels = await asyncio.wait_for(
                session.client.screenshot(x, y, width, height),
                timeout=SCREENSHOT_TIMEOUT,
            )
        except asyncio.TimeoutError:
            raise ToolError(
                f"Screenshot timed out after {SCREENSHOT_TIMEOUT}s. "
                f"The VNC server may be unresponsive. Try again or disconnect."
            )
        except Exception as e:
            raise ToolError(f"Screenshot failed: {e}")

        session.width = session.client.video.width
        session.height = session.client.video.height

    png_bytes = _rgba_to_png(pixels, max_width)
    meta = {
        "session_id": session_id,
        "region": {"x": x, "y": y, "width": width, "height": height},
        "desktop_size": {"width": session.width, "height": session.height},
        "png_size_bytes": len(png_bytes),
    }
    if max_width > 0 and width > max_width:
        ratio = max_width / width
        meta["scaled_width"] = max_width
        meta["scaled_height"] = int(height * ratio)
        meta["note"] = "Coordinates in other tools use original resolution, not scaled."

    return ToolResult(content=[json.dumps(meta), Image(data=png_bytes, format="png")])


@mcp.tool()
async def vnc_wait_for_screen_change(
    ctx: Context,
    session_id: str,
    timeout: int = WAIT_CHANGE_TIMEOUT,
    x: int | None = None,
    y: int | None = None,
    width: int | None = None,
    height: int | None = None,
) -> str:
    """Wait for the remote screen to change after an interaction.

    Call this after clicking a button, pressing a key, etc. to wait until
    the UI has actually updated before taking the next screenshot.

    Optionally specify a region to watch -- only changes within that region
    will count. Without a region, any screen update triggers a return.

    Does NOT return an image. Call vnc_screenshot after this if you need
    to see the updated screen.

    Args:
        session_id: VNC session to watch.
        timeout: Max seconds to wait (default 5).
        x: Left edge of region to watch (requires y, width, height).
        y: Top edge of region to watch.
        width: Width of region to watch.
        height: Height of region to watch.
    """
    state = get_state(ctx)
    session = state.get(session_id)
    video = session.client.video

    # Must have taken at least one screenshot so video.data is populated
    if video.data is None:
        raise ToolError(
            "No screen data available -- take a vnc_screenshot first "
            "so the framebuffer is populated."
        )

    # Validate region params: all-or-nothing
    region_params = [x, y, width, height]
    has_region = any(p is not None for p in region_params)
    if has_region:
        if any(p is None for p in region_params):
            raise ToolError(
                "Region requires all four params: x, y, width, height "
                f"(got x={x}, y={y}, width={width}, height={height})"
            )
        _validate_region(session, x, y, width, height)

    # Snapshot the "before" state
    serial_before = video.serial
    if has_region:
        before_pixels = video.as_rgba(x, y, width, height).copy()

    changed = False
    updates_received = 0
    t_start = time.monotonic()

    async with session.lock:
        # Request incremental update from the server
        video.refresh()
        await session.client.drain()

        while True:
            remaining = timeout - (time.monotonic() - t_start)
            if remaining <= 0:
                break

            try:
                update_type = await asyncio.wait_for(
                    session.client.read(), timeout=remaining
                )
            except asyncio.TimeoutError:
                break
            except (asyncio.IncompleteReadError, ConnectionError) as e:
                raise ToolError(f"VNC connection lost during wait: {e}")

            # Import here to avoid top-level dependency on asyncvnc2 internals
            from asyncvnc2 import UpdateType

            if update_type == UpdateType.VIDEO:
                updates_received += 1
                if video.serial != serial_before:
                    if has_region:
                        current_pixels = video.as_rgba(x, y, width, height)
                        if not np.array_equal(before_pixels, current_pixels):
                            changed = True
                            break
                        # Serial changed but our region didn't -- keep waiting
                    else:
                        changed = True
                        break
            # BELL and CLIPBOARD updates are consumed silently

        # Update cached dimensions
        session.width = video.width
        session.height = video.height

    elapsed = round(time.monotonic() - t_start, 3)
    result = {
        "changed": changed,
        "elapsed_seconds": elapsed,
        "serial_before": serial_before,
        "serial_after": video.serial,
        "updates_received": updates_received,
    }
    if has_region:
        result["region"] = {"x": x, "y": y, "width": width, "height": height}
    if not changed:
        result["note"] = (
            f"No change detected within {timeout}s. The UI may need more "
            "time, or the interaction may not have triggered a visual update."
        )

    return json.dumps(result)
