"""VM discovery and inspection tools (vm_list, vm_info, vm_screenshot, vm_screenshot_api)."""

from __future__ import annotations

import io
import logging
import re
import tempfile
import time
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Annotated

import anyio
from fastmcp.exceptions import ToolError
from fastmcp.tools import tool
from fastmcp.utilities.types import Image
from PIL import Image as PilImage
from pydantic import Field

from mcp_vm_blackbox._helpers import SCRATCH_DIR, _load_vboxapi, _resolve, _ssh_run, _vbox
from mcp_vm_blackbox.models import ScaleInfo, ScreenshotResult, ScreenshotUnchangedResult, VmInfoResult, VmListEntry
from mcp_vm_blackbox.vm_interaction import update_marks_cache
from mcp_vm_blackbox.vm_registry import default_registry

_log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Perceptual hash deduplication (Layer 2)
# ---------------------------------------------------------------------------

#: Hamming distance threshold for declaring two frames identical.
_HASH_THRESHOLD: int = 8

#: Target dimensions for the perceptual hash thumbnail.
_HASH_SIZE: int = 8


def _compute_perceptual_hash(png_bytes: bytes) -> bytes:
    """Resize image to 8x8 greyscale and return raw pixel bytes as hash.

    Args:
        png_bytes: Raw PNG image bytes.

    Returns:
        64-byte sequence of greyscale pixel values used as a perceptual hash.
    """
    with PilImage.open(io.BytesIO(png_bytes)) as img:
        thumb = img.convert("L").resize((_HASH_SIZE, _HASH_SIZE), PilImage.Resampling.LANCZOS)
        return thumb.tobytes()


def _hamming_distance(a: bytes, b: bytes) -> int:
    """Count the number of differing bytes between two equal-length byte strings.

    Args:
        a: First byte sequence.
        b: Second byte sequence.

    Returns:
        Number of positions where ``a`` and ``b`` differ.
    """
    return sum(x != y for x, y in zip(a, b, strict=False))


class PerceptualHashStore:
    """In-process store mapping VM name to the most recent perceptual hash.

    The hash is computed from the full-resolution PNG (pre-optimisation) so
    that dedup comparisons are consistent with evidence stored on disk.

    Attributes:
        _hashes: Maps vm_name -> (monotonic_time, hash_bytes).
    """

    def __init__(self) -> None:
        """Initialise an empty hash store."""
        self._hashes: dict[str, tuple[float, bytes]] = {}

    def store(self, vm_name: str, png_bytes: bytes) -> bytes:
        """Compute and store the perceptual hash for *vm_name*.

        Args:
            vm_name: VirtualBox VM name used as the store key.
            png_bytes: Raw PNG bytes of the current frame.

        Returns:
            The computed hash bytes (64 bytes for an 8x8 greyscale thumb).
        """
        hash_bytes = _compute_perceptual_hash(png_bytes)
        self._hashes[vm_name] = (time.monotonic(), hash_bytes)
        return hash_bytes

    def check_unchanged(self, vm_name: str, png_bytes: bytes, threshold: int = _HASH_THRESHOLD) -> tuple[bool, float]:
        """Compare current frame hash against the stored baseline.

        Args:
            vm_name: VirtualBox VM name.
            png_bytes: Raw PNG bytes of the current frame.
            threshold: Maximum Hamming distance to consider frames identical.

        Returns:
            ``(unchanged, elapsed_seconds)`` where *unchanged* is ``True``
            when the frame matches the stored baseline within *threshold*,
            and *elapsed_seconds* is the time since the last stored capture
            (``0.0`` when no baseline exists).
        """
        if vm_name not in self._hashes:
            return False, 0.0
        stored_time, stored_hash = self._hashes[vm_name]
        current_hash = _compute_perceptual_hash(png_bytes)
        distance = _hamming_distance(current_hash, stored_hash)
        elapsed = time.monotonic() - stored_time
        return distance <= threshold, elapsed


#: Module-level singleton shared across all tool calls in a process.
_hash_store = PerceptualHashStore()


# ---------------------------------------------------------------------------
# JPEG optimisation helper (Layer 1)
# ---------------------------------------------------------------------------


def _to_jpeg(png_bytes: bytes, quality: int = 85, max_size: tuple[int, int] = (1024, 768)) -> tuple[bytes, ScaleInfo]:
    """Downscale a PNG to XGA and encode as JPEG.

    The aspect ratio is preserved.  The returned bytes are suitable for
    context delivery; the original full-resolution PNG must be saved to the
    evidence store separately.

    Args:
        png_bytes: Raw PNG bytes from the VM screenshot capture.
        quality: JPEG quality (0-95).  Defaults to 85.
        max_size: Maximum (width, height) for the output image.

    Returns:
        Tuple of (JPEG-encoded bytes, ScaleInfo describing the downscale ratio).
        When native dimensions are zero or image already fits within max_size,
        scale_x and scale_y are 1.0.
    """
    with PilImage.open(io.BytesIO(png_bytes)) as img:
        native_width, native_height = img.size
        img.thumbnail(max_size, PilImage.Resampling.LANCZOS)
        output_width, output_height = img.size
        buf = io.BytesIO()
        img.convert("RGB").save(buf, format="JPEG", quality=quality)

    if native_width == 0 or native_height == 0:
        scale_x = 1.0
        scale_y = 1.0
    else:
        scale_x = output_width / native_width
        scale_y = output_height / native_height

    scale_info = ScaleInfo(
        scale_x=scale_x,
        scale_y=scale_y,
        native_width=native_width,
        native_height=native_height,
        output_width=output_width,
        output_height=output_height,
    )
    return buf.getvalue(), scale_info


def _parse_vm_list(all_output: str, running_output: str) -> list[VmListEntry]:
    """Parse ``VBoxManage list vms`` and ``list runningvms`` output into records.

    Args:
        all_output: Raw stdout from ``VBoxManage list vms``.
        running_output: Raw stdout from ``VBoxManage list runningvms``.

    Returns:
        List of VmListEntry models.
    """
    pattern = re.compile(r'"([^"]+)"\s+\{([^}]+)\}')
    running_names = {m.group(1) for m in pattern.finditer(running_output)}
    return [
        VmListEntry(name=m.group(1), uuid=m.group(2), running=m.group(1) in running_names)
        for m in pattern.finditer(all_output)
    ]


def _screenshot_via_api(vm_name: str, out: Path) -> ScreenshotResult | None:
    """Attempt a screenshot using the vboxapi Python SDK.

    Args:
        vm_name: VirtualBox VM name.
        out: Destination path for the saved PNG.

    Returns:
        Result dict with ``saved_to``, ``width``, ``height`` on success,
        or ``None`` when vboxapi is not available (caller should fall back).
    """
    try:
        vbox_mgr_cls = _load_vboxapi()
    except ToolError:
        return None

    mgr = vbox_mgr_cls(None, None)
    try:
        vbox = mgr.getVirtualBox()
        machine = vbox.findMachine(vm_name)
        session = mgr.openMachineSession(machine, fPermitSharing=True)
        try:
            display = session.console.display
            w, h, _bpp, _xo, _yo, _gs = display.getScreenResolution(0)
            png_bytes: bytes = bytes(display.takeScreenShotToArray(0, w, h, mgr.constants.BitmapFormat_PNG))
        finally:
            mgr.closeMachineSession(session)
    except Exception as exc:
        raise RuntimeError(f"vboxapi screenshot failed: {exc}") from exc

    _ = out.write_bytes(png_bytes)
    return ScreenshotResult(saved_to=str(out), width=w, height=h)


async def _screenshot_via_subprocess(vm_name: str, target: str, out: Path) -> ScreenshotResult:
    """Take a screenshot via ``VBoxManage screenshotpng`` async fallback.

    Args:
        vm_name: VirtualBox VM name.
        target: Named target or raw SSH string.
        out: Destination path for the saved PNG.

    Returns:
        Result dict with ``saved_to``, ``width`` (0), ``height`` (0), or
        ``{"error": ...}`` when the VM is not found.
    """
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        tmp_path = tmp.name
    try:
        cfg = _resolve(target)
        dispatch = _vbox(cfg, vm_name)
        png_bytes = await dispatch.screenshot(tmp_path)
    except RuntimeError as exc:
        return ScreenshotResult(
            saved_to="", width=0, height=0, error=f"VM '{vm_name}' not found. Run vm_list to see available VMs. ({exc})"
        )
    finally:
        await anyio.Path(tmp_path).unlink(missing_ok=True)

    await anyio.Path(out).write_bytes(png_bytes)
    return ScreenshotResult(saved_to=str(out), width=0, height=0)


async def _apply_som_overlay(jpeg_bytes: bytes, scale_info: ScaleInfo, vm_name: str, target: str) -> bytes:
    """Detect UI elements and stamp Set-of-Marks labels on a JPEG image.

    Called by :func:`vm_screenshot` when ``som=True``.  Extracted into a
    helper to keep ``vm_screenshot`` within the local-variable limit.

    When element detection returns nothing (VM unreachable, not Windows,
    PowerShell error) the original *jpeg_bytes* are returned unchanged.

    Args:
        jpeg_bytes: JPEG-encoded bytes from the downscale step.
        scale_info: Scale metadata from :func:`_to_jpeg`.
        vm_name: VirtualBox VM name (used for element detection and cache key).
        target: Named target or raw SSH string (forwarded to detect_elements).

    Returns:
        JPEG bytes -- either stamped with marks or the original when no
        elements were detected.
    """
    from mcp_vm_blackbox.som_overlay import detect_elements, stamp_marks  # noqa: PLC0415

    elements = await detect_elements(vm_name=vm_name, target=target)
    if not elements:
        _log.debug("vm_screenshot: SoM requested but no elements detected for %s", vm_name)
        return jpeg_bytes

    with PilImage.open(io.BytesIO(jpeg_bytes)) as jpeg_img:
        stamped_img, legend_draft = stamp_marks(jpeg_img.copy(), elements, scale_info)
        buf = io.BytesIO()
        stamped_img.convert("RGB").save(buf, format="JPEG", quality=85)
        stamped_bytes = buf.getvalue()

    # stamp_marks leaves vm_name as "" -- patch it with the actual VM name.
    legend = legend_draft.model_copy(update={"vm_name": vm_name})
    update_marks_cache(vm_name, legend)
    _log.debug("vm_screenshot: SoM stamped %d marks for %s", len(legend.marks), vm_name)
    return stamped_bytes


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@tool
async def vm_list(
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> list[VmListEntry]:
    """List all VirtualBox VMs on the target host with their running state.

    Args:
        target: Named target (``"local"``, ``"ci"``) or raw SSH string
            (``"user@host"``).

    Returns:
        List of VmListEntry with name, uuid, running.
    """
    cfg = _resolve(target)
    all_result = await _ssh_run(cfg, ["VBoxManage", "list", "vms"])
    running_result = await _ssh_run(cfg, ["VBoxManage", "list", "runningvms"])
    entries = _parse_vm_list(all_result.stdout, running_result.stdout)
    try:
        await default_registry.populate(entries)
    except Exception:  # noqa: BLE001
        _log.warning("Failed to update VM registry", exc_info=True)
    return entries


@tool
async def vm_screenshot(
    vm_name: Annotated[str, Field(description="VirtualBox VM name (use vm_list to discover names).")],
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
    output_path: Annotated[
        str | None,
        Field(
            description="Optional explicit destination path for the saved PNG. "
            "When None, a unique timestamped path is generated inside scratch/screenshots/."
        ),
    ] = None,
    force: Annotated[
        bool, Field(description="When True, bypass perceptual-hash deduplication and always return a fresh JPEG image.")
    ] = False,
    som: Annotated[
        bool,
        Field(
            description="When True, stamp numbered Set-of-Marks labels on the returned "
            "image. Elements are detected via UIAutomation (Windows VMs only). "
            "The marks legend is cached so vm_mouse_click can resolve mark IDs to "
            "native coordinates. The evidence PNG is always saved pre-marks; marks "
            "appear only in the context image."
        ),
    ] = False,
) -> Image | ScreenshotUnchangedResult:
    """Take a screenshot of a VM and return a downscaled JPEG for context efficiency.

    Saves the full-resolution PNG to ``scratch/screenshots/<vm_name>-<timestamp>-<uid>.png``
    for evidence preservation, then applies optimisation layers before returning:

    * **Layer 1** -- JPEG downscale: resizes to at most 1024x768 (XGA) at JPEG
      quality 85, reducing context bytes by ~95 % compared with a raw PNG.
    * **Layer 2** -- Perceptual dedup: compares an 8x8 greyscale hash of the
      current frame against the stored baseline.  When the Hamming distance is
      <= 8 (out of 64 bytes) the screen is considered unchanged and a
      :class:`~mcp_vm_blackbox.models.ScreenshotUnchangedResult` is returned
      instead of an image, eliminating image tokens entirely.
    * **SoM overlay** (optional, ``som=True``) -- after JPEG downscale, detects
      interactive UI elements via UIAutomation and stamps numbered labels on the
      image.  The marks legend is cached so ``vm_mouse_click`` can resolve mark
      IDs to native coordinates.  The evidence PNG is always pre-marks.

    Pass ``force=True`` to bypass Layer 2 and always receive a fresh image.

    Args:
        vm_name: VirtualBox VM name (use :func:`vm_list` to discover names).
        target: Named target (``"local"``, ``"ci"``) or raw SSH string
            (``"user@host"``).
        output_path: Optional explicit destination path for the saved PNG.
            When ``None``, a unique timestamped path is generated inside
            ``scratch/screenshots/`` so concurrent callers never clobber each
            other.
        force: When ``True``, skip the dedup check and always return a JPEG.
        som: When ``True``, stamp Set-of-Marks numbered labels on the returned
            image.  Marks are applied **after** JPEG downscale and are never
            written to the evidence store.  When no elements are detected the
            image is returned without marks.

    Returns:
        :class:`~fastmcp.utilities.types.Image` with JPEG bytes when the
        screen has changed, or
        :class:`~mcp_vm_blackbox.models.ScreenshotUnchangedResult` when the
        screen is unchanged and ``force`` is ``False``.
        Raises :class:`RuntimeError` on capture failure.
    """
    cfg = _resolve(target)
    dispatch = _vbox(cfg, vm_name)

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        tmp_path = tmp.name

    try:
        png_bytes = await dispatch.screenshot(tmp_path)
    except RuntimeError as exc:
        raise RuntimeError(f"VM '{vm_name}' not found. Run vm_list to see available VMs. ({exc})") from exc
    finally:
        await anyio.Path(tmp_path).unlink(missing_ok=True)

    # --- Layer 2: perceptual dedup (pre-optimisation hash for consistency) ---
    # Hash is always computed on raw PNG -- before marks and before JPEG encode.
    if not force:
        unchanged, elapsed = _hash_store.check_unchanged(vm_name, png_bytes)
        if unchanged:
            return ScreenshotUnchangedResult(
                message=f"Screen unchanged since last capture ({elapsed:.1f}s ago)",
                last_capture_seconds_ago=elapsed,
                vm_name=vm_name,
            )

    # Store hash from full-res PNG (pre-optimisation) for next comparison.
    _hash_store.store(vm_name, png_bytes)

    # --- Evidence: always save full-resolution pre-marks PNG to disk ---
    if output_path is not None:
        out = anyio.Path(output_path)
        await out.parent.mkdir(parents=True, exist_ok=True)
    else:
        await anyio.Path(SCRATCH_DIR).mkdir(parents=True, exist_ok=True)
        ts = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
        uid = uuid.uuid4().hex[:8]
        out = anyio.Path(SCRATCH_DIR) / f"{vm_name}-{ts}-{uid}.png"
    await out.write_bytes(png_bytes)

    # --- Layer 1: JPEG downscale for context delivery ---
    jpeg_bytes, scale_info = _to_jpeg(png_bytes)

    # --- SoM overlay (optional): stamp marks on the downscaled JPEG ---
    if som:
        jpeg_bytes = await _apply_som_overlay(jpeg_bytes, scale_info, vm_name, target)

    return Image(data=jpeg_bytes, format="jpeg")


@tool
async def vm_info(
    vm_name: Annotated[str, Field(description="VirtualBox VM name.")],
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> VmInfoResult:
    """Return VM hardware and state info (memory, CPU, VRAM, guest additions).

    Args:
        vm_name: VirtualBox VM name.
        target: Named target or raw SSH string.

    Returns:
        VmInfoResult with memory_mb, cpus, vram_mb, guest_additions, state.
        Returns VmInfoResult with error set when the VM is not found.
    """
    cfg = _resolve(target)
    result = await _ssh_run(cfg, ["VBoxManage", "showvminfo", vm_name, "--machinereadable"])
    if result.returncode != 0:
        return VmInfoResult(error=f"VM '{vm_name}' not found. Run vm_list to see available VMs.")

    info: dict[str, str] = {}
    for line in result.stdout.splitlines():
        if "=" in line:
            k, _, v = line.partition("=")
            info[k.strip()] = v.strip().strip('"')

    return VmInfoResult(
        memory_mb=info.get("memory", "unknown"),
        cpus=info.get("cpus", "unknown"),
        vram_mb=info.get("vram", "unknown"),
        state=info.get("VMState", "unknown"),
        guest_additions=info.get("GuestAdditionsVersion", "not installed"),
    )


@tool
async def vm_screenshot_api(
    vm_name: Annotated[str, Field(description="VirtualBox VM name (use vm_list to discover names).")],
    output_path: Annotated[
        str | None,
        Field(
            description="Optional explicit destination path for the saved PNG. When None, a unique timestamped path is generated inside scratch/screenshots/."
        ),
    ] = None,
    target: Annotated[
        str,
        Field(description='Named target or raw SSH string. Only "local" is supported (vboxapi is local XPCOM only).'),
    ] = "local",
) -> ScreenshotResult:
    """Take a screenshot via the VirtualBox Python API -- no subprocess, no temp file race condition.

    Uses ``display.takeScreenShotToArray`` directly.  Falls back to async
    ``VBoxManage screenshotpng`` if vboxapi is unavailable.

    Only works for ``target="local"`` because vboxapi uses local XPCOM.

    Args:
        vm_name: VirtualBox VM name (use :func:`vm_list` to discover names).
        output_path: Optional explicit destination path for the saved PNG.
            When ``None``, a unique timestamped path is generated inside
            ``scratch/screenshots/`` so concurrent callers never clobber
            each other.
        target: Named target or raw SSH string.  Must be ``"local"``.

    Returns:
        Dict with ``saved_to`` (local file path string), ``width`` (int),
        and ``height`` (int).  Returns ``{"error": ...}`` on failure.

    Raises:
        ToolError: When ``target`` is not ``"local"``.
    """
    if target != "local":
        raise ToolError("vm_screenshot_api only supports target='local' (vboxapi uses local XPCOM)")

    if output_path is not None:
        out = Path(output_path)
        out.parent.mkdir(parents=True, exist_ok=True)
    else:
        SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
        ts = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
        uid = uuid.uuid4().hex[:8]
        out = SCRATCH_DIR / f"{vm_name}-{ts}-{uid}.png"

    result = _screenshot_via_api(vm_name, out)
    if result is not None:
        return result
    return await _screenshot_via_subprocess(vm_name, target, out)
