from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from pathlib import Path

import httpx
import pytest
from fastapi.testclient import TestClient

from cloud.auth import digest_token
from cloud.control_config import CloudControlConfig
from cloud.sdk.client import CloudClient
from cloud.scheduler import TaskConstraints
from cloud_api.app import create_app
from core.models import Scene, utc_now
from device.manager import DeviceManager
from driver.base import Driver
from host_agent.assignment import AssignmentExecutionResult, AssignmentExecutor
from host_agent.client import HostAgentClient, StaleLeaseError
from host_agent.config import HostAgentConfig
from host_agent.execution import ExecutionFactories
from host_agent.heartbeat import HeartbeatSynchronizer
from host_agent.lease import ActiveAssignmentRunner
from host_agent.processor import AssignmentProcessor
from runtime.executor import Executor, default_tool_registry
from runtime.planner import PlannedStep, Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from workflow.store import WorkflowStore


class FakeDriver(Driver):
    def __init__(self) -> None:
        self.calls: list[tuple[str, tuple[object, ...]]] = []

    def connect(self) -> None:
        self.calls.append(("connect", ()))

    def disconnect(self) -> None:
        return None

    def screenshot(self) -> bytes:
        return b""

    def tap(self, x: float, y: float) -> None:
        self.calls.append(("tap", (x, y)))

    def long_press(self, x: float, y: float, duration_ms: int = 1200) -> None:
        return None

    def swipe(
        self,
        start_x: float,
        start_y: float,
        end_x: float,
        end_y: float,
        duration_ms: int = 500,
    ) -> None:
        return None

    def input(self, text: str) -> None:
        return None

    def launch(self, app_id: str) -> None:
        return None

    def terminate(self, app_id: str) -> None:
        return None

    def tree(self):
        return None

    def home(self) -> None:
        return None

    def lock(self) -> None:
        return None

    def unlock(self) -> None:
        return None


def _config(host_id: str) -> HostAgentConfig:
    return HostAgentConfig(
        control_plane_url="http://control.test",
        host_id=host_id,
        token=f"token-{host_id}",
        poll_timeout_seconds=0.01,
        retry_backoff_seconds=0.001,
        max_retry_backoff_seconds=0.001,
    )


@asynccontextmanager
async def _control_plane(
    database_path: Path,
    *host_ids: str,
    lease_duration_seconds: float = 30,
) -> AsyncIterator[object]:
    app = create_app(
        config=CloudControlConfig(
            environment="test",
            database_url=f"sqlite:///{database_path.as_posix()}",
            scheduler_interval_seconds=60,
            lease_reaper_interval_seconds=60,
            lease_duration_seconds=lease_duration_seconds,
        )
    )
    async with app.router.lifespan_context(app):
        for host_id in host_ids:
            app.state.cloud_services.repository.enroll_host(
                host_id=host_id,
                agent_instance_id=f"agent-{host_id}",
                credential_digest=digest_token(f"token-{host_id}"),
                enrollment_token_digest=None,
                display_name=host_id,
                enrolled_at=utc_now(),
            )
        yield app


@asynccontextmanager
async def _host_client(
    app,
    host_id: str,
    *,
    request_paths: list[str] | None = None,
) -> AsyncIterator[HostAgentClient]:
    transport: httpx.AsyncBaseTransport = httpx.ASGITransport(app=app)
    if request_paths is not None:
        transport = _RecordingTransport(transport, request_paths)
    async with httpx.AsyncClient(
        transport=transport,
        base_url="http://control.test",
    ) as http_client:
        yield HostAgentClient(_config(host_id), http_client=http_client)


class _RecordingTransport(httpx.AsyncBaseTransport):
    def __init__(
        self,
        transport: httpx.AsyncBaseTransport,
        paths: list[str],
    ) -> None:
        self.transport = transport
        self.paths = paths

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        self.paths.append(request.url.path)
        return await self.transport.handle_async_request(request)

    async def aclose(self) -> None:
        await self.transport.aclose()


async def _sync_fake_device(
    app,
    client: HostAgentClient,
    host_id: str,
    device_id: str,
    *,
    driver_type: str = "wda",
) -> FakeDriver:
    app.state.cloud_services.repository.enroll_device(
        device_id=device_id,
        host_id=host_id,
        local_device_id=device_id,
        driver_type=driver_type,
        name=device_id,
        capability_tags=[],
        enrolled_at=utc_now(),
    )
    driver = FakeDriver()
    manager = DeviceManager()
    manager.register_device(
        device_id,
        lambda: driver,
        driver_type=driver_type,
        status="idle",
    )
    await HeartbeatSynchronizer(manager, client, _config(host_id)).sync_once()
    return driver


class _SuccessfulExecution:
    async def run(self, assignment):
        return AssignmentExecutionResult(
            status="done",
            metadata={"fake_driver": assignment.device_id},
        )

    def request_stop(self) -> None:
        return None


class _AsyncTestClient:
    def __init__(self, client: TestClient) -> None:
        self.client = client

    async def request(self, method: str, path: str, **kwargs) -> httpx.Response:
        kwargs.pop("timeout", None)
        return await asyncio.to_thread(self.client.request, method, path, **kwargs)


class _TapPlanner(Planner):
    def plan(self, *, goal, scene, context):
        if context.step_results:
            return []
        return [PlannedStep(action="tap", description="tap", args={"x": 2, "y": 3})]

    def goal_reached(self, *, goal, scene, context):
        return bool(context.step_results)


class _FailingPlanner(Planner):
    def plan(self, *, goal, scene, context):
        raise RuntimeError("planner unavailable")


def _assignment_processor(
    tmp_path: Path,
    client: HostAgentClient,
    manager: DeviceManager,
    *,
    planner: Planner,
) -> AssignmentProcessor:
    scene = Scene(width=10, height=20, elements=[])

    def task_runner_factory() -> TaskRunner:
        return TaskRunner(
            planner=planner,
            executor=Executor(tools=default_tool_registry(manager=manager)),
            observer=lambda device_id: scene,
            screenshot_provider=lambda device_id: b"",
            config=TaskRunnerConfig(max_steps=3),
        )

    workflow_store = WorkflowStore(tmp_path / "workflows.sqlite3")
    factories = ExecutionFactories(
        task_runner_factory=task_runner_factory,
        workflow_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_store=workflow_store,
    )
    executor = AssignmentExecutor(factories)
    return AssignmentProcessor(client, ActiveAssignmentRunner(client, executor))


def test_one_host_executes_assignment_through_outbound_protocol(tmp_path) -> None:
    async def scenario() -> None:
        paths: list[str] = []
        async with _control_plane(tmp_path / "one-host.sqlite3", "host-a") as app:
            async with _host_client(app, "host-a", request_paths=paths) as client:
                await _sync_fake_device(app, client, "host-a", "device-a")
                task_id = app.state.cloud_services.scheduler.submit(
                    goal="open settings"
                )
                app.state.cloud_services.scheduler.assign()
                assignment = await client.claim()
                assert assignment is not None

                result = await AssignmentProcessor(
                    client,
                    _SuccessfulExecution(),
                ).process(assignment)

                task = app.state.cloud_services.repository.get_task(task_id)
                assert result.report_status == "recorded"
                assert task.status == "done"
                assert task.terminal_result == {"fake_driver": "device-a"}
        assert all(path.startswith("/internal/v1/") for path in paths)

    asyncio.run(scenario())


def test_nat_style_host_requires_only_outbound_requests(tmp_path) -> None:
    async def scenario() -> None:
        paths: list[str] = []
        async with _control_plane(tmp_path / "outbound-only.sqlite3", "host-a") as app:
            async with _host_client(app, "host-a", request_paths=paths) as client:
                await _sync_fake_device(app, client, "host-a", "device-a")
                assert await client.claim() is None

        assert paths == [
            "/internal/v1/hosts/host-a/heartbeat",
            "/internal/v1/hosts/host-a/assignments/claim",
        ]

    asyncio.run(scenario())


def test_multiple_hosts_claim_only_their_matching_devices(tmp_path) -> None:
    async def scenario() -> None:
        async with _control_plane(
            tmp_path / "multiple-hosts.sqlite3",
            "host-a",
            "host-b",
        ) as app:
            async with (
                _host_client(app, "host-a") as client_a,
                _host_client(app, "host-b") as client_b,
            ):
                await _sync_fake_device(app, client_a, "host-a", "device-a")
                await _sync_fake_device(
                    app,
                    client_b,
                    "host-b",
                    "device-b",
                    driver_type="appium",
                )
                task_a = app.state.cloud_services.scheduler.submit(
                    goal="ios task",
                    constraints=TaskConstraints(driver_type="wda"),
                )
                task_b = app.state.cloud_services.scheduler.submit(
                    goal="android task",
                    constraints=TaskConstraints(driver_type="appium"),
                )
                app.state.cloud_services.scheduler.assign()

                assignment_a = await client_a.claim()
                assignment_b = await client_b.claim()

                assert assignment_a is not None and assignment_a.task_id == task_a
                assert assignment_b is not None and assignment_b.task_id == task_b
                assert assignment_a.host_id == "host-a"
                assert assignment_b.host_id == "host-b"

    asyncio.run(scenario())


def test_control_plane_restart_preserves_dispatched_assignment(tmp_path) -> None:
    async def scenario() -> None:
        database_path = tmp_path / "control-restart.sqlite3"
        assignment = None
        async with _control_plane(database_path, "host-a") as first_app:
            async with _host_client(first_app, "host-a") as client:
                await _sync_fake_device(first_app, client, "host-a", "device-a")
                task_id = first_app.state.cloud_services.scheduler.submit(goal="resume")
                first_app.state.cloud_services.scheduler.assign()
                assignment = await client.claim()
                assert assignment is not None

        async with _control_plane(database_path, "host-a") as restarted_app:
            async with _host_client(restarted_app, "host-a") as client:
                assert (await client.renew(assignment)).status == "renewed"
                await client.report_result(assignment, status="done")
                task = restarted_app.state.cloud_services.repository.get_task(task_id)
                assert task.status == "done"

    asyncio.run(scenario())


def test_host_agent_restart_reuses_active_lease(tmp_path) -> None:
    async def scenario() -> None:
        async with _control_plane(tmp_path / "agent-restart.sqlite3", "host-a") as app:
            async with _host_client(app, "host-a") as first_client:
                await _sync_fake_device(app, first_client, "host-a", "device-a")
                task_id = app.state.cloud_services.scheduler.submit(goal="resume host")
                app.state.cloud_services.scheduler.assign()
                assignment = await first_client.claim()
                assert assignment is not None

            async with _host_client(app, "host-a") as restarted_client:
                assert (await restarted_client.renew(assignment)).status == "renewed"
                await restarted_client.report_result(assignment, status="done")

            assert (
                app.state.cloud_services.repository.get_task(task_id).status == "done"
            )

    asyncio.run(scenario())


def test_lease_loss_rejects_stale_host_result(tmp_path) -> None:
    async def scenario() -> None:
        async with _control_plane(
            tmp_path / "lease-loss.sqlite3",
            "host-a",
            lease_duration_seconds=0.2,
        ) as app:
            async with _host_client(app, "host-a") as client:
                await _sync_fake_device(app, client, "host-a", "device-a")
                task_id = app.state.cloud_services.scheduler.submit(goal="expire")
                app.state.cloud_services.scheduler.assign()
                assignment = await client.claim()
                assert assignment is not None
                await asyncio.sleep(0.25)
                app.state.cloud_services.repository.reap_expired_leases(
                    now=datetime.now(UTC),
                    max_attempts=3,
                )

                with pytest.raises(StaleLeaseError):
                    await client.report_result(assignment, status="done")

                task = app.state.cloud_services.repository.get_task(task_id)
                assert task.status == "queued"
                assert task.terminal_result is None

    asyncio.run(scenario())


def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) -> None:
    app = create_app(
        config=CloudControlConfig(
            environment="test",
            database_url=f"sqlite:///{(tmp_path / 'sdk-e2e.sqlite3').as_posix()}",
            scheduler_interval_seconds=60,
            lease_reaper_interval_seconds=60,
            lease_duration_seconds=30,
        )
    )
    driver = FakeDriver()
    manager = DeviceManager()
    manager.register_device("device-a", lambda: driver, status="idle")

    with TestClient(app) as http_client:
        app.state.cloud_services.user_auth_service.create_user(
            username="operator",
            display_name="Operator",
            role="admin",
            password="correct-horse-battery-staple",
            must_change_password=False,
        )
        login = http_client.post(
            "/v1/auth/login",
            json={"username": "operator", "password": "correct-horse-battery-staple"},
        )
        assert login.status_code == 200
        app.state.cloud_services.repository.enroll_host(
            host_id="host-a",
            agent_instance_id="agent-host-a",
            credential_digest=digest_token("token-host-a"),
            enrollment_token_digest=None,
            display_name="host-a",
            enrolled_at=utc_now(),
        )
        app.state.cloud_services.repository.enroll_device(
            device_id="device-a",
            host_id="host-a",
            local_device_id="device-a",
            driver_type="wda",
            name="device-a",
            capability_tags=[],
            enrolled_at=utc_now(),
        )
        cloud_client = CloudClient(
            "http://testserver",
            http_client=http_client,
        )

        async def scenario() -> None:
            host_client = HostAgentClient(
                _config("host-a"),
                http_client=_AsyncTestClient(http_client),  # type: ignore[arg-type]
            )
            heartbeat = HeartbeatSynchronizer(manager, host_client, _config("host-a"))
            heartbeat.connect_devices()
            await heartbeat.sync_once()

            successful_task_id = cloud_client.submit_task(goal="tap screen")["task_id"]
            app.state.cloud_services.scheduler.assign()
            successful_assignment = await host_client.claim()
            assert successful_assignment is not None
            await _assignment_processor(
                tmp_path / "success",
                host_client,
                manager,
                planner=_TapPlanner(),
            ).process(successful_assignment)

            successful_status = cloud_client.get_task_status(successful_task_id)
            assert successful_status["status"] == "done"
            assert ("tap", (2, 3)) in driver.calls

            failed_task_id = cloud_client.submit_task(goal="fail planning")["task_id"]
            app.state.cloud_services.scheduler.assign()
            failed_assignment = await host_client.claim()
            assert failed_assignment is not None
            await _assignment_processor(
                tmp_path / "failure",
                host_client,
                manager,
                planner=_FailingPlanner(),
            ).process(failed_assignment)

            failed_status = cloud_client.get_task_status(failed_task_id)
            assert failed_status["status"] == "failed"
            assert failed_status["failure_reason"] == (
                "RuntimeError: planner unavailable"
            )

        asyncio.run(scenario())