# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Portions of this file are derived from mobile-use (https://github.com/minitap-ai/mobile-use)
# Copyright 2025-2026 Minitap, Inc. Licensed under the Apache License 2.0.

"""UIAutomator2 client for Android device screen data retrieval.

Provides an alternative to the Maestro-based screen API with direct device
access.
Handles Maestro blocker detection and removal before connecting.
"""

import base64
from io import BytesIO
import re as _re
import subprocess
import time
from typing import TYPE_CHECKING
import xml.etree.ElementTree as ET

from PIL import Image
from pydantic import BaseModel
import uiautomator2 as u2

from artemis.runtime.adb_endpoint import adb_command
from artemis.runtime.awake_service import ensure_device_awake
from artemis.utils.logger import get_logger

if TYPE_CHECKING:
    from uiautomator2 import Device

logger = get_logger(__name__)


MAESTRO_PACKAGE = "dev.mobile.maestro"


class _CyFunctionDetectorMeta(type):
    def __instancecheck__(self, instance):
        name = type(instance).__name__
        return (
            name
            in (
                "cyfunction",
                "cython_function_or_method",
                "builtin_function_or_method",
            )
            or "cyfunction" in name.lower()
        )


class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta):
    pass


class UIAutomatorScreenData(BaseModel):
    """Screen data response from UIAutomator2."""

    model_config = {"ignored_types": (CyFunctionDetector,)}

    base64: str
    hierarchy_xml: str
    elements: list[dict]
    width: int
    height: int


def _parse_hierarchy_xml_to_elements(hierarchy_xml: str) -> list[dict]:
    """Parse uiautomator2 XML hierarchy into a flat list of element dictionaries.

    The output format matches the existing screen API format with attributes
    like:
    - resource-id
    - text
    - content-desc
    - bounds (as string "[x1,y1][x2,y2]")
    - class
    - package
    - checkable, checked, clickable, enabled, focusable, focused
    - scrollable, long-clickable, password, selected

    Args:
        hierarchy_xml: XML string from uiautomator2.dump_hierarchy()

    Returns:
        Flat list of element dictionaries
    """
    elements: list[dict] = []

    try:
        root = ET.fromstring(hierarchy_xml)
    except ET.ParseError as e:
        logger.error(f"Failed to parse hierarchy XML: {e}")
        return elements

    def _extract_element(node: ET.Element) -> None:
        """Recursively extract elements from XML nodes."""
        element: dict = {}

        for attr_name, attr_value in node.attrib.items():
            if attr_name == "resource-id":
                element["resource-id"] = attr_value
            elif attr_name == "text":
                element["text"] = attr_value
            elif attr_name == "content-desc":
                element["content-desc"] = attr_value
                element["accessibilityText"] = attr_value
            elif attr_name == "bounds":
                element["bounds"] = attr_value

                match = _re.match(r"\[(-?\d+),(-?\d+)\]\[(-?\d+),(-?\d+)\]", attr_value)
                if match:
                    x1, y1, x2, y2 = map(int, match.groups())
                    element["parsed_bounds"] = {
                        "left": x1,
                        "top": y1,
                        "right": x2,
                        "bottom": y2,
                    }
            elif attr_name == "class":
                element["class"] = attr_value
            elif attr_name == "package":
                element["package"] = attr_value
            elif attr_name in (
                "checkable",
                "checked",
                "clickable",
                "enabled",
                "focusable",
                "focused",
                "scrollable",
                "long-clickable",
                "password",
                "selected",
            ):
                element[attr_name] = attr_value
            else:
                element[attr_name] = attr_value

        if element:
            elements.append(element)

        for child in node:
            _extract_element(child)

    _extract_element(root)

    return elements


def _pil_to_base64(img: Image.Image, format: str = "JPEG", quality: int = 80) -> str:
    """Convert PIL Image to base64 string with high-quality compression."""
    buffer = BytesIO()
    if format.upper() in ("JPEG", "JPG"):
        format = "JPEG"
        if img.mode in ("RGBA", "P"):
            img = img.convert("RGB")
        img.save(buffer, format=format, quality=quality, optimize=True)
    else:
        img.save(buffer, format=format)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _is_package_installed(device_id: str, pkg: str) -> bool:
    """Check if a package is installed on the device."""
    try:
        result = subprocess.run(
            adb_command(["-s", device_id, "shell", "pm", "list", "packages"]),
            stdin=subprocess.DEVNULL,
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            logger.warning(f"Failed to list packages: {result.stderr}")
            return False

        lines = result.stdout.splitlines()
        target = f"package:{pkg}"
        return target in lines
    except subprocess.TimeoutExpired:
        logger.warning("Timeout checking installed packages")
        return False
    except (OSError, subprocess.SubprocessError) as e:
        logger.warning(f"Error checking installed packages: {e}")
        return False


def _uninstall_package(device_id: str, pkg: str) -> bool:
    """Uninstall a package from the device for the current user."""
    try:
        result = subprocess.run(
            adb_command(
                [
                    "-s",
                    device_id,
                    "shell",
                    "pm",
                    "uninstall",
                    "--user",
                    "0",
                    pkg,
                ]
            ),
            stdin=subprocess.DEVNULL,
            capture_output=True,
            text=True,
            timeout=30,
        )
        if result.returncode == 0:
            logger.info(f"Successfully uninstalled {pkg}")
            return True
        else:
            logger.warning(f"Failed to uninstall {pkg}: {result.stderr}")
            return False
    except subprocess.TimeoutExpired:
        logger.warning(f"Timeout uninstalling {pkg}")
        return False
    except Exception as e:
        logger.warning(f"Error uninstalling {pkg}: {e}")
        return False


def _ensure_maestro_not_installed(device_id: str) -> None:
    """Check if Maestro is installed and uninstall it.

    Maestro conflicts with uiautomator2 - if Maestro's package is installed,
    u2.connect() will fail. This function ensures Maestro is removed before
    attempting to connect.
    """
    if _is_package_installed(device_id, MAESTRO_PACKAGE):
        logger.warning(
            f"Maestro ({MAESTRO_PACKAGE}) detected - uninstalling to enable UIAutomator2..."
        )
        _uninstall_package(device_id, MAESTRO_PACKAGE)


class UIAutomatorConnectionError(RuntimeError):
    """Terminal connection failure with device context and the original cause."""

    def __init__(self, device_id: str, attempts: int, cause: Exception) -> None:
        # The library appends its whole output buffer on every readiness poll.
        # Keep the original exception as __cause__, but do not repeat those
        # identical startup lines dozens of times in the user-facing failure.
        lines = [line for arg in cause.args for line in str(arg).splitlines()]
        detail = "\n".join(dict.fromkeys(lines)) or str(cause)
        super().__init__(
            f"UIAutomator2 连接失败（设备：{device_id}；已尝试 {attempts} 次）。"
            f"\n底层原因：{type(cause).__name__}: {detail}"
        )


class _ManagedUIAutomatorDevice(u2.Device):
    """Release a failed launch before the constructor can lose its handle.

    UIAutomator2 registers exit cleanup only after startup succeeds. Its
    partially constructed Device is not returned when readiness fails, so
    the caller cannot stop that launch. Intercept the startup boundary here
    instead of modifying the installed library or killing device processes.
    """

    def start_uiautomator(self) -> None:
        try:
            super().start_uiautomator()
        except BaseException:
            # _process is set only for a server launched by this instance;
            # attaching to a server already on the device leaves it None.
            if self._process is not None:
                try:
                    super().stop_uiautomator(wait=False)
                except Exception as cleanup_error:
                    # Cleanup must never hide the startup error or interrupt.
                    logger.warning(
                        f"Failed to stop this UIAutomator2 launch on {self._dev.serial}: "
                        f"{cleanup_error}"
                    )
            raise


class UIAutomatorClient:
    """UIAutomator2 client for Android screen data retrieval.

    This client uses uiautomator2 library for direct device communication,
    providing faster hierarchy and screenshot retrieval compared to Maestro.

    Important: Maestro must not be installed on the device as it conflicts
    with uiautomator2. This client automatically handles Maestro removal.
    """

    def __init__(self, device_id: str):
        """Initialize the UIAutomator client.

        Args:
            device_id: The Android device serial number (e.g., "emulator-5554")
        """
        self._device_id = device_id
        self._device: Device | None = None
        self._awake_strategy: str | None = None

    def _ensure_connected(self) -> "Device":
        """Ensure connection to the device, handling Maestro blocker.

        Returns:
            Connected uiautomator2 Device instance
        """
        if self._device is not None:
            try:
                # Quick check if connection is still alive
                self._device.info
                return self._device
            except Exception:
                logger.warning("UIAutomator2 connection lost, reconnecting...")
                self._device = None

        # Ensure Maestro is not blocking us
        _ensure_maestro_not_installed(self._device_id)

        # Enroll the device in the containing Artemis service lifetime before
        # UIAutomator2 reads the screen. Client disconnects do not release it.
        if self._awake_strategy is None:
            self._awake_strategy = ensure_device_awake(self._device_id)

        # Retry connections while ADB recovers from a reboot or USB disconnect.
        logger.info(f"Connecting UIAutomator2 to device: {self._device_id}")
        last_error: Exception | None = None
        for attempt in range(3):
            if attempt:
                time.sleep(1.0 * attempt)
            try:
                self._device = _ManagedUIAutomatorDevice(self._device_id)
                break
            except Exception as exc:
                last_error = exc
                logger.warning(
                    f"UIAutomator2 connect attempt {attempt + 1}/3 to"
                    f" {self._device_id} failed: {exc}"
                )
        else:
            self._awake_strategy = None
            raise UIAutomatorConnectionError(
                self._device_id, attempt + 1, last_error
            ) from last_error
        logger.info("UIAutomator2 connected successfully")
        return self._device

    def connect(self) -> None:
        """Connect and activate this session's awake strategy without an action."""
        self._ensure_connected()

    def press_key(self, key: str):
        """Press a key on the device.

        Args:
            key: Key to press (e.g., "home", "back", "enter"...)
        """
        device = self._ensure_connected()
        return device.press(key=key)

    def send_text(self, text: str) -> None:
        """Send text input to the device using FastInputIME.

        This method supports special characters (e.g., 'ö') that ADB shell
        input text cannot handle. The original IME is restored after input.

        Args:
            text: The text to input
        """

        device = self._ensure_connected()
        device.set_fastinput_ime(True)
        # Allow IME to initialize
        time.sleep(0.3)
        try:
            device.send_keys(text)
            # Give FastInputIME time to process the broadcast and commit text
            # before switching it off and killing it.
            time.sleep(0.5)
        finally:
            device.set_fastinput_ime(False)

    def clear_text(self) -> bool:
        """Clear text of the currently focused element using UIAutomator2.

        Returns True on success, False on failure.
        """
        try:
            device = self._ensure_connected()
            device(focused=True).clear_text()
            return True
        except Exception as e:
            logger.error(f"UIAutomator2 clear_text failed: {e}")
            return False

    def set_clipboard(self, text: str) -> bool:
        """Set text to system clipboard via UIAutomator2.

        Args:
            text: The string to set in the device clipboard.

        Returns:
            True if successful, False otherwise.
        """
        try:
            device = self._ensure_connected()
            device.set_clipboard(text)
            return True
        except Exception as e:
            logger.debug(f"UIAutomator2 set_clipboard failed: {e}")
            return False

    def get_hierarchy(self) -> str:
        """Get the UI hierarchy XML from the device.

        Returns:
            Compressed UI hierarchy XML string
        """
        device = self._ensure_connected()
        return device.dump_hierarchy(compressed=True)

    def get_screenshot(self) -> Image.Image | None:
        """Capture a screenshot from the device.

        Uses adb screencap for more reliable capture of all layers (including
        popups), falling back to UIAutomator2 if it fails.

        Returns:
            PIL Image or None if capture failed
        """
        try:
            result = subprocess.run(
                adb_command(["-s", self._device_id, "exec-out", "screencap", "-p"]),
                stdin=subprocess.DEVNULL,
                capture_output=True,
                timeout=10,
                check=True,
            )
            return Image.open(BytesIO(result.stdout))
        except (OSError, subprocess.SubprocessError) as e:
            logger.warning(
                f"Failed to capture screenshot via adb: {e}, falling back to uiautomator2"
            )
            device = self._ensure_connected()
            return device.screenshot()

    def get_screenshot_base64(self) -> str | None:
        """Capture a screenshot and return as base64 string.

        Returns:
            Base64 encoded PNG screenshot or None if capture failed
        """
        screenshot = self.get_screenshot()
        if screenshot is None:
            return None
        return _pil_to_base64(screenshot, format="JPEG")

    def get_screen_data(self) -> UIAutomatorScreenData:
        """Get complete screen data including screenshot and hierarchy.

        Returns:
            UIAutomatorScreenData with screenshot, hierarchy, elements, and
            dimensions
        """
        device = self._ensure_connected()

        # Get screenshot
        screenshot = self.get_screenshot()
        if screenshot is None:
            raise RuntimeError("Failed to capture screenshot")

        # Get hierarchy
        hierarchy_xml = device.dump_hierarchy(compressed=True)

        # Parse XML to flat elements list
        elements = _parse_hierarchy_xml_to_elements(hierarchy_xml)

        return UIAutomatorScreenData(
            base64=_pil_to_base64(screenshot, format="JPEG"),
            hierarchy_xml=hierarchy_xml,
            elements=elements,
            width=screenshot.width,
            height=screenshot.height,
        )

    def disconnect(self, stop_server: bool = False) -> None:
        """Disconnect this client without ending the host's awake lifetime.

        ``stop_server`` also kills uiautomator2's device server. Its UiAutomation
        connection suppresses every accessibility service (the Artemis helper
        included) for as long as it lives, so a caller that prefers the helper
        must release it; the pure UIAutomator2 backend keeps it for fast reconnects.
        """
        device = self._device
        self._awake_strategy = None
        self._device = None
        if stop_server and device is not None:
            self.stop_server(device)
        logger.info("UIAutomator2 client disconnected")

    def stop_server(self, device: "Device | None" = None) -> bool:
        """Kill uiautomator2's on-device server (see :meth:`disconnect`)."""
        device = device or self._device
        if device is None:
            return False
        try:
            device.stop_uiautomator()
            logger.info(f"UIAutomator2 device server stopped on {self._device_id}")
            return True
        except (
            OSError,
            RuntimeError,
            subprocess.SubprocessError,
            u2.exceptions.DeviceError,
        ) as exc:
            logger.debug(f"Stopping the UIAutomator2 server on {self._device_id} failed: {exc}")
            return False


def get_client(device_id: str) -> UIAutomatorClient:
    """Factory function to create a UIAutomatorClient.

    Args:
        device_id: The Android device serial number

    Returns:
        UIAutomatorClient instance
    """
    return UIAutomatorClient(device_id=device_id)