# 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.

"""Unit tests for Unified Base Device Driver and implementations."""

import asyncio
import base64
import threading
from io import BytesIO
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest
from PIL import Image

from artemis.drivers.base import KeyCode, ScreenData, SwipeDirection
from artemis.drivers.mock.mock_driver import MockDeviceDriver
from artemis.drivers.android.adb_driver import AndroidAdbDriver
from artemis.config.agent import ScreenshotConfig
from artemis.utils.screenshot_limit import limit_screenshot_long_edge


_ONE_PIXEL_PNG = base64.b64encode(
    base64.b64decode(
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
    )
).decode("ascii")


def _valid_png():
    buffer = BytesIO()
    Image.new("RGB", (64, 64), "white").save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("ascii")


_VALID_PNG = _valid_png()


@pytest.mark.asyncio
async def test_android_driver_rejects_placeholder_without_returning_stale_screen():
    client = MagicMock()
    client.get_screen_data.return_value = SimpleNamespace(base64=_ONE_PIXEL_PNG)
    driver = AndroidAdbDriver("test-device", MagicMock(), ui_adb_client=client)
    with pytest.raises(RuntimeError, match="占位图"):
        await driver.get_screen_data(skip_settling=True)


@pytest.mark.asyncio
async def test_screen_capture_timeout_does_not_start_overlapping_sdk_calls(monkeypatch):
    released = threading.Event()
    client = MagicMock()

    def blocked_capture():
        released.wait(2)
        return SimpleNamespace(base64=_VALID_PNG, width=64, height=64, elements=[])

    client.get_screen_data.side_effect = blocked_capture
    driver = AndroidAdbDriver("test-device", MagicMock(), ui_adb_client=client)
    monkeypatch.setattr(
        "artemis.drivers.android.adb_driver.settings.ARTEMIS_OBSERVATION_TIMEOUT_SECONDS", 0.02
    )
    try:
        for _ in range(3):
            with pytest.raises(TimeoutError, match="观测超时"):
                await driver.get_screen_data(skip_settling=True)
        assert client.get_screen_data.call_count == 1
    finally:
        released.set()
        await asyncio.wait_for(driver._screen_capture_task, 2)


def test_screenshot_limit_keeps_small_and_invalid_images_unchanged():
    small_image = base64.b64decode(_ONE_PIXEL_PNG)
    assert limit_screenshot_long_edge(small_image, 2048) is small_image
    invalid_image = b"not an image"
    assert limit_screenshot_long_edge(invalid_image, 2048) is invalid_image


@pytest.mark.asyncio
async def test_mock_driver_actions():
    """Verify MockDeviceDriver records and performs all core actions."""
    driver = MockDeviceDriver(device_id="test-mock", width=1080, height=2400)
    assert driver.device_id == "test-mock"
    assert driver.screen_size == (1080, 2400)

    # 1. Screen data
    screen_data = await driver.get_screen_data()
    assert isinstance(screen_data, ScreenData)
    assert screen_data.width == 1080
    assert screen_data.height == 2400
    assert screen_data.screenshot_base64 is not None

    # 2. Tap
    assert await driver.tap(500, 1000) is True
    assert driver.action_history[-1]["action"] == "tap"
    assert driver.action_history[-1]["x"] == 500

    # 3. Swipe
    assert await driver.swipe(100, 200, 300, 400) is True
    assert driver.action_history[-1]["action"] == "swipe"

    # 4. Swipe direction
    assert await driver.swipe_direction(SwipeDirection.UP) is True
    assert driver.action_history[-1]["direction"] == "up"

    # 5. Input text
    assert await driver.input_text("Hello Artemis") is True
    assert driver.action_history[-1]["text"] == "Hello Artemis"

    # 6. Press key
    assert await driver.press_key(KeyCode.HOME) is True
    assert driver.action_history[-1]["key"] == "home"

    # 7. App lifecycle
    assert await driver.launch_app("com.android.settings") is True
    assert await driver.get_current_package() == "com.android.settings"
    assert await driver.stop_app("com.android.settings") is True


@pytest.mark.asyncio
async def test_android_driver_with_mock_adb():
    """Verify AndroidAdbDriver formats shell commands properly."""
    mock_adb_client = MagicMock()
    mock_adb_device = MagicMock()
    mock_adb_client.device.return_value = mock_adb_device

    driver = AndroidAdbDriver(
        device_id="emulator-5554",
        adb_client=mock_adb_client,
        width=1080,
        height=2400,
    )

    # Test tap
    await driver.tap(540, 1200)
    mock_adb_device.shell.assert_called_with("input tap 540 1200", timeout=10.0)

    # Test multi-tap
    await driver.tap(540, 1200, times=3, delay_ms=100)
    expected_multi = "input tap 540 1200 && sleep 0.100 && input tap 540 1200 && sleep 0.100 && input tap 540 1200"
    mock_adb_device.shell.assert_called_with(expected_multi, timeout=10.0)

    # Test press key
    await driver.press_key(KeyCode.BACK)
    mock_adb_device.shell.assert_called_with("input keyevent 4", timeout=10.0)


@pytest.mark.asyncio
async def test_android_driver_preserves_elements_from_combined_screen_data():
    """The combined uiautomator2 response is the authoritative live snapshot."""
    mock_adb_client = MagicMock()
    mock_ui_client = MagicMock()
    expected_elements = [
        {
            "text": "",
            "class": "android.widget.FrameLayout",
            "bounds": "[0,0][1080,2400]",
            "clickable": "false",
            "focusable": "false",
        },
        {
            "text": "Start",
            "class": "android.widget.Button",
            "bounds": "[10,20][100,80]",
        },
    ]
    mock_ui_client.get_screen_data.return_value = SimpleNamespace(
        base64=_VALID_PNG,
        hierarchy_xml="<hierarchy />",
        elements=expected_elements,
        width=1080,
        height=2400,
    )

    driver = AndroidAdbDriver(
        device_id="emulator-5554",
        adb_client=mock_adb_client,
        ui_adb_client=mock_ui_client,
    )

    screen_data = await driver.get_screen_data(skip_settling=True)

    assert screen_data.ui_elements == [expected_elements[1]]
    assert screen_data.ui_hierarchy_xml == "<hierarchy />"
    mock_ui_client.get_hierarchy.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.parametrize("enabled", [True, False])
async def test_android_screenshot_limit_preserves_native_coordinates(enabled):
    image = Image.new("RGB", (1200, 2600), "white")
    buffer = BytesIO()
    image.save(buffer, format="JPEG")
    original = buffer.getvalue()
    elements = [{"text": "目标", "bounds": "[100,200][300,400]"}]
    ui_client = MagicMock()
    ui_client.get_screen_data.return_value = SimpleNamespace(
        base64=base64.b64encode(original).decode("ascii"),
        hierarchy_xml="<hierarchy />",
        elements=elements,
        width=1200,
        height=2600,
    )
    driver = AndroidAdbDriver("emulator-5554", MagicMock(), ui_adb_client=ui_client)
    driver._screenshot_config = ScreenshotConfig(enabled=enabled, max_long_edge=2048)

    screen = await driver.get_screen_data(skip_settling=True)

    with Image.open(BytesIO(screen.screenshot_bytes)) as result:
        assert result.size == ((945, 2048) if enabled else (1200, 2600))
    assert screen.width == 1200 and screen.height == 2600
    assert screen.ui_elements == elements
    assert screen.screenshot_base64 == base64.b64encode(screen.screenshot_bytes).decode("ascii")
    if not enabled:
        assert screen.screenshot_bytes == original


@pytest.mark.asyncio
async def test_android_driver_parses_xml_hierarchy_fallback():
    """XML-only clients remain supported without discarding their hierarchy."""
    mock_adb_client = MagicMock()
    mock_ui_client = MagicMock()
    mock_ui_client.get_screen_data.return_value = SimpleNamespace(
        base64=_VALID_PNG,
        width=1080,
        height=2400,
    )
    mock_ui_client.get_hierarchy.return_value = (
        '<hierarchy><node text="Settings" class="android.widget.TextView" '
        'bounds="[0,0][100,100]" /></hierarchy>'
    )

    driver = AndroidAdbDriver(
        device_id="emulator-5554",
        adb_client=mock_adb_client,
        ui_adb_client=mock_ui_client,
    )

    screen_data = await driver.get_screen_data(skip_settling=True)

    assert screen_data.ui_hierarchy_xml is not None
    assert any(element.get("text") == "Settings" for element in screen_data.ui_elements)


@pytest.mark.asyncio
async def test_find_element_prefers_fresh_bounds_over_stale_center():
    """Layout transitions must not leave dynamic element taps at old coordinates."""
    driver = MockDeviceDriver(device_id="test-mock", width=1080, height=2400)
    driver.get_screen_data = AsyncMock(
        return_value=ScreenData(
            screenshot_bytes=base64.b64decode(_ONE_PIXEL_PNG),
            screenshot_base64=_VALID_PNG,
            ui_elements=[
                {
                    "resource_id": "app:id/digit_4",
                    "bounds": "[10,1700][270,1940]",
                    "parsed_bounds": {
                        "left": 10,
                        "top": 1700,
                        "right": 270,
                        "bottom": 1940,
                    },
                    "center": [140, 1600],
                }
            ],
            width=1080,
            height=2400,
        )
    )

    _, center, error = await driver.find_element(resource_id="app:id/digit_4")

    assert error is None
    assert center == [140, 1820]