"""Set-of-Marks (SoM) overlay module.

Provides UIAutomation-based element detection and Pillow-based mark rendering
for the screenshot efficiency pipeline.  All SoM logic is encapsulated here so
that vm_inspection.py and viewport_tools.py can apply marks without knowing the
detection or rendering implementation.

Architecture notes:
- detect_elements calls vm_powershell directly (it is an async coroutine).
- stamp_marks operates on an in-memory PIL Image; callers are responsible for
  re-encoding to JPEG after stamping.
- Coordinates in SoMMarksLegend / MarkEntry are always in **native** VM
  resolution (pre-downscale) so they can be forwarded directly to vm_mouse_click.
- detect_elements NEVER raises — it returns an empty list on any failure.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from datetime import UTC, datetime

from PIL import Image as PilImage, ImageDraw, ImageFont

from mcp_vm_blackbox.models import MarkEntry, ScaleInfo, SoMMarksLegend

log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SOM_FONT_SIZE: int = 14
SOM_LABEL_BG_COLOR: tuple[int, int, int] = (255, 0, 0)
SOM_LABEL_FG_COLOR: tuple[int, int, int] = (255, 255, 255)
SOM_MAX_ELEMENTS: int = 50

#: Padding around the text inside the label background rectangle (pixels).
_LABEL_PADDING: int = 2


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class DetectedElement:
    """A UI element detected via Windows UI Automation.

    Attributes:
        name: Human-readable element name as reported by UIAutomation.
        control_type: UIAutomation ControlType string (e.g. 'Button', 'Edit').
        bounding_rect: Element bounds in native VM coordinates:
            (left, top, right, bottom).
    """

    name: str
    control_type: str
    bounding_rect: tuple[int, int, int, int]


# ---------------------------------------------------------------------------
# PowerShell script builder
# ---------------------------------------------------------------------------


def build_uiautomation_script() -> str:
    """Return a PowerShell script that enumerates interactive UI elements.

    The script uses the built-in System.Windows.Automation namespace (part of
    .NET Framework on Windows — no extra packages required).  It queries the
    foreground window's element tree, filters for elements that are both
    visible and have a non-empty name, then serialises them as a JSON array.

    Each element in the JSON array has the following fields:
    - ``Name`` (string): the accessibility name of the element.
    - ``ControlType`` (string): the local-name portion of the ControlType
      property (e.g. "Button", "Edit", "CheckBox").
    - ``Left``, ``Top``, ``Right``, ``Bottom`` (int): bounding rectangle in
      screen coordinates (native VM resolution).

    Returns:
        PowerShell source text suitable for passing directly to vm_powershell.
    """
    return r"""
Add-Type -AssemblyName UIAutomationClient
Add-Type -AssemblyName UIAutomationTypes

$root = [System.Windows.Automation.AutomationElement]::RootElement
$scope = [System.Windows.Automation.TreeScope]::Descendants
$cond = [System.Windows.Automation.Condition]::TrueCondition

$interactive = @(
    [System.Windows.Automation.ControlType]::Button,
    [System.Windows.Automation.ControlType]::CheckBox,
    [System.Windows.Automation.ControlType]::ComboBox,
    [System.Windows.Automation.ControlType]::Edit,
    [System.Windows.Automation.ControlType]::Hyperlink,
    [System.Windows.Automation.ControlType]::List,
    [System.Windows.Automation.ControlType]::ListItem,
    [System.Windows.Automation.ControlType]::MenuItem,
    [System.Windows.Automation.ControlType]::RadioButton,
    [System.Windows.Automation.ControlType]::Slider,
    [System.Windows.Automation.ControlType]::Spinner,
    [System.Windows.Automation.ControlType]::Tab,
    [System.Windows.Automation.ControlType]::TabItem,
    [System.Windows.Automation.ControlType]::Text,
    [System.Windows.Automation.ControlType]::ToolBar,
    [System.Windows.Automation.ControlType]::Tree,
    [System.Windows.Automation.ControlType]::TreeItem
)

$elements = $root.FindAll($scope, $cond)
$results = New-Object System.Collections.Generic.List[PSObject]

foreach ($el in $elements) {
    try {
        $name = $el.GetCurrentPropertyValue([System.Windows.Automation.AutomationElement]::NameProperty)
        if (-not $name -or $name.Trim() -eq '') { continue }

        $ctRaw = $el.GetCurrentPropertyValue([System.Windows.Automation.AutomationElement]::ControlTypeProperty)
        $ct = $ctRaw.LocalizedControlType
        if (-not $ct) { $ct = $ctRaw.Id.ToString() }

        $rect = $el.GetCurrentPropertyValue([System.Windows.Automation.AutomationElement]::BoundingRectangleProperty)
        $left   = [int]$rect.Left
        $top    = [int]$rect.Top
        $right  = [int]$rect.Right
        $bottom = [int]$rect.Bottom

        if ($right -le $left -or $bottom -le $top) { continue }

        $interactiveType = $interactive | Where-Object { $_.Id -eq $ctRaw.Id }
        if (-not $interactiveType) { continue }

        $obj = [PSCustomObject]@{
            Name         = [string]$name
            ControlType  = [string]$ct
            Left         = $left
            Top          = $top
            Right        = $right
            Bottom       = $bottom
        }
        $results.Add($obj)
    } catch {
        # Skip elements that throw (e.g. elements from other sessions)
        continue
    }
}

$results | ConvertTo-Json -Compress -Depth 1
""".strip()


# ---------------------------------------------------------------------------
# Element detection
# ---------------------------------------------------------------------------


async def detect_elements(vm_name: str, target: str = "local") -> list[DetectedElement]:
    """Detect interactive UI elements on the VM screen via UIAutomation.

    Calls ``vm_powershell`` with a hardcoded UIAutomation PowerShell script and
    parses the JSON response.  On any failure (WinRM error, JSON parse failure,
    empty result), returns an empty list rather than raising.

    Args:
        vm_name: VirtualBox VM name (forwarded to vm_powershell for context).
        target: Named target or raw SSH string (forwarded to vm_powershell).

    Returns:
        Ordered list of detected elements (up to SOM_MAX_ELEMENTS).  Empty on
        failure or when no interactive elements are visible.
    """
    # Import here to avoid circular imports; vm_interaction imports models, not us.
    import asyncssh  # noqa: PLC0415
    from requests.exceptions import RequestException  # noqa: PLC0415

    from mcp_vm_blackbox.vm_interaction import vm_powershell  # noqa: PLC0415

    script = build_uiautomation_script()
    try:
        result = await vm_powershell(vm_name=vm_name, script=script, target=target)
    except (FileNotFoundError, ValueError, OSError, RuntimeError, asyncssh.Error, RequestException):
        log.debug("detect_elements: vm_powershell call failed for %s", vm_name, exc_info=True)
        return []

    if result.exit_code != 0 or not result.stdout.strip():
        log.debug(
            "detect_elements: PowerShell exited %d for %s; stderr=%r",
            result.exit_code,
            vm_name,
            result.stderr[:200] if result.stderr else "",
        )
        return []

    try:
        raw = json.loads(result.stdout)
    except json.JSONDecodeError:
        log.debug("detect_elements: JSON parse failed for %s; stdout=%r", vm_name, result.stdout[:200])
        return []

    # PowerShell ConvertTo-Json returns a single object (not an array) when the
    # collection has exactly one item.  Normalise to list.
    if isinstance(raw, dict):
        raw = [raw]
    if not isinstance(raw, list):
        log.debug("detect_elements: unexpected JSON shape for %s: %r", vm_name, type(raw))
        return []

    elements: list[DetectedElement] = []
    for item in raw:
        if len(elements) >= SOM_MAX_ELEMENTS:
            break
        try:
            name = str(item["Name"]).strip()
            if not name:
                continue
            ct = str(item["ControlType"]).strip() or "Unknown"
            left = int(item["Left"])
            top = int(item["Top"])
            right = int(item["Right"])
            bottom = int(item["Bottom"])
            if right <= left or bottom <= top:
                continue
            elements.append(DetectedElement(name=name, control_type=ct, bounding_rect=(left, top, right, bottom)))
        except (KeyError, TypeError, ValueError):
            log.debug("detect_elements: skipping malformed element item: %r", item)
            continue

    return elements


# ---------------------------------------------------------------------------
# Mark rendering helpers
# ---------------------------------------------------------------------------


def _load_font(font_size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    """Load Pillow's built-in default font at *font_size* pixels.

    Pillow >= 10.1 supports ``ImageFont.load_default(size=N)``.  Older builds
    do not accept a size argument; fall back silently.

    Args:
        font_size: Desired font size in pixels.

    Returns:
        A Pillow font object suitable for ``ImageDraw.text()``.
    """
    try:
        return ImageFont.load_default(size=font_size)
    except TypeError:
        return ImageFont.load_default()


def _draw_mark_label(
    draw: ImageDraw.ImageDraw,
    font: ImageFont.FreeTypeFont | ImageFont.ImageFont,
    label: str,
    left_s: int,
    top_s: int,
    img_size: tuple[int, int],
) -> None:
    """Draw a single numbered label rectangle onto *draw*.

    Args:
        draw: An ImageDraw context for the target image.
        font: Font to use for the label text.
        label: Text string to draw (e.g. "1", "42").
        left_s: Left edge of the label in scaled (output) coordinates.
        top_s: Top edge of the label in scaled (output) coordinates.
        img_size: ``(width, height)`` of the target image for bounds clamping.
    """
    bbox = draw.textbbox((0, 0), label, font=font)
    rect_x1 = min(left_s + (bbox[2] - bbox[0]) + _LABEL_PADDING * 2, img_size[0])
    rect_y1 = min(top_s + (bbox[3] - bbox[1]) + _LABEL_PADDING * 2, img_size[1])
    draw.rectangle((left_s, top_s, rect_x1, rect_y1), fill=SOM_LABEL_BG_COLOR)
    draw.text((left_s + _LABEL_PADDING, top_s + _LABEL_PADDING), label, fill=SOM_LABEL_FG_COLOR, font=font)


# ---------------------------------------------------------------------------
# Mark rendering — public API
# ---------------------------------------------------------------------------


def stamp_marks(
    image: PilImage.Image, elements: list[DetectedElement], scale_info: ScaleInfo, font_size: int = SOM_FONT_SIZE
) -> tuple[PilImage.Image, SoMMarksLegend]:
    """Stamp numbered mark labels onto *image* at each detected element position.

    Labels are positioned at the top-left corner of each element's bounding
    rect after scaling from native to output (downscaled) coordinates.  Each
    label is a filled red rectangle with a white number drawn on it.

    The input *image* is modified in place and also returned for convenience.

    Args:
        image: PIL Image (RGB or RGBA) to draw marks on.  Modified in place.
        elements: Detected UI elements (capped at SOM_MAX_ELEMENTS).
        scale_info: Scale factors and dimensions from the JPEG downscale step.
        font_size: Point size for the Pillow built-in font.

    Returns:
        A (image, legend) tuple where:
        - ``image`` is the modified PIL Image with marks drawn on it.
        - ``legend`` is a SoMMarksLegend with one MarkEntry per stamped element,
          storing native-resolution coordinates.
    """
    draw = ImageDraw.Draw(image)
    font = _load_font(font_size)
    marks: list[MarkEntry] = []

    for idx, elem in enumerate(elements[:SOM_MAX_ELEMENTS], start=1):
        left_n, top_n, right_n, bottom_n = elem.bounding_rect
        _draw_mark_label(
            draw, font, str(idx), int(left_n * scale_info.scale_x), int(top_n * scale_info.scale_y), image.size
        )
        marks.append(
            MarkEntry(
                mark_id=idx,
                name=elem.name,
                control_type=elem.control_type,
                native_x=(left_n + right_n) // 2,
                native_y=(top_n + bottom_n) // 2,
                bounding_rect=elem.bounding_rect,
            )
        )

    legend = SoMMarksLegend(
        vm_name="",  # Caller fills in vm_name after the fact.
        marks=marks,
        detection_method="uiautomation",
        timestamp=datetime.now(UTC).isoformat(),
    )
    return image, legend
