"""UIAutomator startup recovery, using the real library with isolated device IO."""

from dataclasses import dataclass, field
import json
import threading
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest
from uiautomator2 import base, core

from artemis.clients import ui_automator_client
from artemis.clients.accessibility_client import HelperUnavailable
from artemis.clients.screen_client_factory import FallbackScreenClient


@dataclass
class Clock:
    now: float = 0.0

    def time(self):
        return self.now

    def sleep(self, seconds):
        self.now += seconds


class DeviceSocket:
    def __init__(self, mode):
        self.mode = mode
        self.closed = threading.Event()
        self.started = threading.Event()

    def recv(self, _size):
        if not self.started.is_set():
            self.started.set()
            output = b"[server] INFO: [UiAutomator2Server] Starting Server\n"
            if self.mode == "conflict":
                output += (
                    b"java.lang.IllegalStateException: UiAutomationService already registered!"
                )
            return output
        self.closed.wait()
        return b""


class DeviceConnection:
    def __init__(self, mode):
        self.conn = DeviceSocket(mode)
        self.close_calls = 0
        self.close_error = None

    def close(self):
        self.close_calls += 1
        if self.close_error:
            raise self.close_error
        self.conn.closed.set()


@dataclass
class DeviceServer:
    modes: list[str] = field(default_factory=lambda: ["not_ready"])
    connections: list[DeviceConnection] = field(default_factory=list)
    previous_closed_at_launch: list[bool] = field(default_factory=list)
    clock: Clock = field(default_factory=Clock)
    cleanup_callbacks: MagicMock = field(default_factory=MagicMock)
    close_error: Exception | None = None
    probe_error: BaseException | None = None

    def shell(self, command, stream=False):
        assert stream is True
        assert "com.wetest.uia2.Main" in command
        self.previous_closed_at_launch.append(all(c.conn.closed.is_set() for c in self.connections))
        index = min(len(self.connections), len(self.modes) - 1)
        connection = DeviceConnection(self.modes[index])
        connection.close_error = self.close_error
        self.connections.append(connection)
        return connection

    def http_request(self, _device, _port, method, path, *args, **kwargs):
        if self.connections:
            latest = self.connections[-1]
            assert latest.conn.started.wait(timeout=2)
            if self.probe_error is not None:
                raise self.probe_error
            if latest.conn.mode == "ready" and not latest.conn.closed.is_set():
                if method == "POST":
                    return core.HTTPResponse(
                        json.dumps({"result": {"displayWidth": 1080}}).encode()
                    )
                return core.HTTPResponse(b"pong")
        raise core.HTTPError("simulated device server is not ready")


@pytest.fixture
def server(monkeypatch):
    simulated = DeviceServer()
    device = SimpleNamespace(serial="isolated-startup-device", shell=simulated.shell)
    monkeypatch.setattr(base._BaseClient, "_wait_for_device", lambda _self: device)
    monkeypatch.setattr(core.BasicUiautomatorServer, "_setup_jar", lambda _self: None)
    monkeypatch.setattr(core, "_http_request", simulated.http_request)
    monkeypatch.setattr(core, "time", simulated.clock)
    monkeypatch.setattr(core, "atexit", SimpleNamespace(register=simulated.cleanup_callbacks))
    monkeypatch.setattr(ui_automator_client, "time", simulated.clock)
    monkeypatch.setattr(ui_automator_client, "_ensure_maestro_not_installed", lambda _s: None)
    monkeypatch.setattr(ui_automator_client, "ensure_device_awake", lambda _s: "simulated")
    yield simulated
    # Close fake sockets even when a regression leaves a real library reader alive.
    for connection in simulated.connections:
        connection.conn.closed.set()


def test_each_failed_launch_is_closed_before_retry_and_final_failure(server):
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(RuntimeError, match="server not ready") as caught:
        client.connect()

    assert len(server.connections) == 3
    assert server.previous_closed_at_launch == [True, True, True]
    assert [c.close_calls for c in server.connections] == [1, 1, 1]
    assert server.clock.now == 93.0
    assert server.cleanup_callbacks.call_count == 0
    assert client._device is None
    assert client._awake_strategy is None
    assert isinstance(caught.value.__cause__, core.LaunchUiAutomationError)
    assert "isolated-startup-device" in str(caught.value)
    assert "3" in str(caught.value)


def test_retry_can_connect_after_failed_launch_was_stopped(server):
    server.modes = ["not_ready", "ready"]
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    client.connect()
    client.connect()  # Cached device info also uses the real library RPC path.

    assert len(server.connections) == 2
    assert server.previous_closed_at_launch == [True, True]
    assert server.connections[0].close_calls == 1
    assert server.connections[1].close_calls == 0
    assert server.cleanup_callbacks.call_count == 1
    client.disconnect(stop_server=True)
    assert server.connections[1].close_calls == 1


def test_later_connection_can_succeed_after_all_previous_attempts_failed(server):
    server.modes = ["not_ready", "not_ready", "not_ready", "ready"]
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(RuntimeError, match="server not ready"):
        client.connect()
    client.connect()

    assert len(server.connections) == 4
    assert server.previous_closed_at_launch == [True, True, True, True]
    assert [c.close_calls for c in server.connections] == [1, 1, 1, 0]
    client.disconnect(stop_server=True)
    assert server.connections[-1].close_calls == 1


def test_auto_fallback_cleans_failed_launches_before_next_attempt(server):
    helper = MagicMock()
    helper.connect.side_effect = HelperUnavailable("simulated helper unavailable")
    client = FallbackScreenClient(
        "isolated-startup-device",
        helper=helper,
        uiautomator_factory=lambda: ui_automator_client.UIAutomatorClient(
            "isolated-startup-device"
        ),
        state_probe=lambda _serial: "device",
    )

    with pytest.raises(RuntimeError, match="server not ready"):
        client.connect()
    client.disconnect()

    assert len(server.connections) == 3
    assert server.previous_closed_at_launch == [True, True, True]
    assert [c.close_calls for c in server.connections] == [1, 1, 1]
    helper.disconnect.assert_called_once()


def test_service_conflict_closes_own_streams_and_preserves_cause(server):
    server.modes = ["conflict"]
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(RuntimeError, match="already registered") as caught:
        client.connect()

    assert isinstance(caught.value.__cause__, core.AccessibilityServiceAlreadyRegisteredError)
    assert len(server.connections) == 3
    assert [c.close_calls for c in server.connections] == [1, 1, 1]


def test_cleanup_error_does_not_replace_original_startup_failure(server):
    server.close_error = OSError("simulated cleanup connection error")
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(RuntimeError, match="server not ready") as caught:
        client.connect()

    assert isinstance(caught.value.__cause__, core.LaunchUiAutomationError)
    assert [c.close_calls for c in server.connections] == [1, 1, 1]


def test_interrupt_during_startup_also_closes_own_stream(server):
    server.probe_error = KeyboardInterrupt()
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(KeyboardInterrupt):
        client.connect()

    assert len(server.connections) == 1
    assert server.connections[0].close_calls == 1


def test_failure_before_launch_does_not_stop_another_device_service(server, monkeypatch):
    failure = OSError("simulated jar transfer failure")

    def fail_before_launch(_self):
        raise failure

    monkeypatch.setattr(core.BasicUiautomatorServer, "_setup_jar", fail_before_launch)
    stop = MagicMock()
    monkeypatch.setattr(core.BasicUiautomatorServer, "stop_uiautomator", stop)
    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")

    with pytest.raises(RuntimeError, match="jar transfer failure") as caught:
        client.connect()

    assert caught.value.__cause__ is failure
    assert server.connections == []
    stop.assert_not_called()


def test_default_helper_success_does_not_start_uiautomator(server):
    factory = MagicMock()
    client = FallbackScreenClient(
        "isolated-startup-device",
        helper=MagicMock(),
        uiautomator_factory=factory,
        state_probe=lambda _serial: "device",
    )

    client.connect()
    client.disconnect()

    factory.assert_not_called()
    assert server.connections == []