"""UI inspection mixin for Android ADB MCP Server.

Provides tools for UI hierarchy inspection and synchronization.
"""

import asyncio
import re
import time
from typing import Any

from fastmcp import Context
from fastmcp.contrib.mcp_mixin import mcp_tool

from .base import ADBBaseMixin


class UIMixin(ADBBaseMixin):
    """Mixin for Android UI inspection.

    Provides tools for:
    - Dumping UI hierarchy (accessibility tree)
    - Waiting for text/elements to appear
    - Finding elements by various attributes
    """

    @mcp_tool()
    async def ui_dump(
        self,
        ctx: Context | None = None,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        """Dump the current UI hierarchy.

        Returns the accessibility tree as XML, showing all visible elements
        with their properties (text, content-description, class, bounds, etc.).

        This is extremely useful for:
        - Finding clickable elements by their text
        - Understanding screen layout without screenshots
        - Locating elements by resource-id or content-description

        Example output element:
        <node text="Settings" class="android.widget.TextView"
              bounds="[0,100][200,150]" clickable="true" />

        Args:
            ctx: MCP context for logging (optional for internal calls)
            device_id: Target device

        Returns:
            UI hierarchy XML and parsed element summary
        """
        if ctx:
            await ctx.info("Dumping UI hierarchy...")

        # Dump UI to temp file on device
        device_path = "/sdcard/window_dump.xml"
        result = await self.run_shell_args(
            ["uiautomator", "dump", device_path], device_id
        )

        if not result.success:
            if ctx:
                await ctx.error(f"UI dump failed: {result.stderr}")
            return {
                "success": False,
                "error": f"Failed to dump UI: {result.stderr}",
            }

        # Read the dump
        cat_result = await self.run_shell_args(["cat", device_path], device_id)

        if not cat_result.success:
            if ctx:
                await ctx.error(f"Failed to read dump: {cat_result.stderr}")
            return {
                "success": False,
                "error": f"Failed to read UI dump: {cat_result.stderr}",
            }

        # Clean up
        await self.run_shell_args(["rm", device_path], device_id)

        xml_content = cat_result.stdout

        # Parse out clickable/important elements for quick reference
        clickable_elements = self._parse_ui_elements(xml_content)

        if ctx:
            await ctx.info(f"Found {len(clickable_elements)} interactive elements")

        return {
            "success": True,
            "xml": xml_content,
            "clickable_elements": clickable_elements,
            "element_count": len(clickable_elements),
        }

    def _parse_ui_elements(self, xml_content: str) -> list[dict[str, Any]]:
        """Parse UI XML to extract clickable/important elements."""
        elements = []

        # Regex to find node elements with their attributes
        node_pattern = re.compile(r"<node\s+([^>]+?)(?:/>|>)", re.DOTALL)
        attr_pattern = re.compile(r'(\w+)="([^"]*)"')

        for match in node_pattern.finditer(xml_content):
            attrs_str = match.group(1)
            attrs = dict(attr_pattern.findall(attrs_str))

            # Only include elements that are interactive or have useful text
            is_clickable = attrs.get("clickable") == "true"
            is_focusable = attrs.get("focusable") == "true"
            has_text = bool(attrs.get("text", "").strip())
            has_desc = bool(attrs.get("content-desc", "").strip())

            if is_clickable or is_focusable or has_text or has_desc:
                element = {
                    "text": attrs.get("text", ""),
                    "content_desc": attrs.get("content-desc", ""),
                    "class": attrs.get("class", ""),
                    "resource_id": attrs.get("resource-id", ""),
                    "clickable": is_clickable,
                    "bounds": attrs.get("bounds", ""),
                }

                # Parse bounds to get center coordinates for tapping
                bounds = attrs.get("bounds", "")
                if bounds:
                    bounds_match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds)
                    if bounds_match:
                        x1, y1, x2, y2 = map(int, bounds_match.groups())
                        element["center"] = {
                            "x": (x1 + x2) // 2,
                            "y": (y1 + y2) // 2,
                        }

                elements.append(element)

        return elements

    @mcp_tool()
    async def ui_find_element(
        self,
        text: str | None = None,
        content_desc: str | None = None,
        resource_id: str | None = None,
        class_name: str | None = None,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        """Find UI elements matching criteria.

        Searches the current UI for elements matching the specified
        attributes. Returns matching elements with their bounds/center
        for interaction.

        Args:
            text: Find elements with this exact text
            content_desc: Find elements with this content-description
            resource_id: Find elements with this resource ID
            class_name: Find elements of this class
            device_id: Target device

        Returns:
            List of matching elements with tap coordinates
        """
        # Get UI dump (internal call, no ctx)
        dump = await self.ui_dump(device_id=device_id)

        if not dump.get("success"):
            return dump

        elements = dump["clickable_elements"]
        matches = []

        for elem in elements:
            match = True

            if text is not None and elem.get("text") != text:
                match = False
            if content_desc is not None and elem.get("content_desc") != content_desc:
                match = False
            if resource_id is not None and resource_id not in elem.get(
                "resource_id", ""
            ):
                match = False
            if class_name is not None and class_name not in elem.get("class", ""):
                match = False

            if match:
                matches.append(elem)

        return {
            "success": True,
            "matches": matches,
            "count": len(matches),
        }

    @mcp_tool()
    async def wait_for_text(
        self,
        text: str,
        timeout_seconds: float = 10.0,
        poll_interval: float = 0.5,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        """Wait for text to appear on screen.

        Polls the UI hierarchy until the specified text is found
        or timeout. Essential for synchronizing automation flows.

        Args:
            text: Text to wait for (case-sensitive, substring match)
            timeout_seconds: Maximum wait time (default 10s)
            poll_interval: Time between polls in seconds (default 0.5s)
            device_id: Target device

        Returns:
            Success status with element info if found, or timeout error
        """
        start_time = time.time()
        attempts = 0

        while (time.time() - start_time) < timeout_seconds:
            attempts += 1

            # Internal call, no ctx
            dump = await self.ui_dump(device_id=device_id)

            if dump.get("success"):
                for elem in dump.get("clickable_elements", []):
                    if text in elem.get("text", "") or text in elem.get(
                        "content_desc", ""
                    ):
                        return {
                            "success": True,
                            "found": True,
                            "element": elem,
                            "wait_time": round(time.time() - start_time, 2),
                            "attempts": attempts,
                        }

            await asyncio.sleep(poll_interval)

        return {
            "success": False,
            "found": False,
            "error": (f"Text '{text}' not found after {timeout_seconds}s"),
            "attempts": attempts,
        }

    @mcp_tool()
    async def wait_for_text_gone(
        self,
        text: str,
        timeout_seconds: float = 10.0,
        poll_interval: float = 0.5,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        """Wait for text to disappear from screen.

        Useful for waiting for loading indicators to finish,
        dialogs to close, etc.

        Args:
            text: Text to wait to disappear
            timeout_seconds: Maximum wait time (default 10s)
            poll_interval: Time between polls (default 0.5s)
            device_id: Target device

        Returns:
            Success when text is no longer visible, or timeout error
        """
        start_time = time.time()
        attempts = 0

        while (time.time() - start_time) < timeout_seconds:
            attempts += 1

            dump = await self.ui_dump(device_id=device_id)

            if dump.get("success"):
                found = False
                for elem in dump.get("clickable_elements", []):
                    if text in elem.get("text", "") or text in elem.get(
                        "content_desc", ""
                    ):
                        found = True
                        break

                if not found:
                    return {
                        "success": True,
                        "gone": True,
                        "wait_time": round(time.time() - start_time, 2),
                        "attempts": attempts,
                    }

            await asyncio.sleep(poll_interval)

        return {
            "success": False,
            "gone": False,
            "error": (f"Text '{text}' still present after {timeout_seconds}s"),
            "attempts": attempts,
        }

    @mcp_tool()
    async def tap_text(
        self,
        text: str,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        """Find element by text and tap it.

        Convenience method that combines ui_find_element + input_tap.
        Finds the first element containing the text and taps its center.

        Args:
            text: Text of element to tap
            device_id: Target device

        Returns:
            Success status with tapped coordinates
        """
        # Find element
        result = await self.ui_find_element(text=text, device_id=device_id)

        if not result.get("success"):
            return result

        matches = result.get("matches", [])
        if not matches:
            # Try content-desc as fallback
            result = await self.ui_find_element(content_desc=text, device_id=device_id)
            matches = result.get("matches", [])

        if not matches:
            return {
                "success": False,
                "error": f"No element found with text '{text}'",
            }

        element = matches[0]
        center = element.get("center")

        if not center:
            return {
                "success": False,
                "error": "Element found but could not determine coordinates",
                "element": element,
            }

        # Tap the center
        tap_result = await self.run_shell_args(
            ["input", "tap", str(center["x"]), str(center["y"])],
            device_id,
        )

        return {
            "success": tap_result.success,
            "action": "tap_text",
            "text": text,
            "coordinates": center,
            "element": element,
            "error": tap_result.stderr if not tap_result.success else None,
        }
