from __future__ import annotations

import asyncio
from datetime import UTC, datetime

import httpx
import pytest

from cloud.internal_api.models import AssignmentModel, DeviceSnapshotModel
from host_agent.client import HostAgentClient, StaleLeaseError
from host_agent.config import HostAgentConfig


def _config(**overrides) -> HostAgentConfig:
    values = {
        "control_plane_url": "https://control.example",
        "host_id": "host-a",
        "token": "host-secret",
        "retry_backoff_seconds": 0.01,
        "max_retry_backoff_seconds": 0.02,
        "max_retry_attempts": 3,
    }
    values.update(overrides)
    return HostAgentConfig(**values)


def _assignment() -> AssignmentModel:
    return AssignmentModel(
        task_id="task-a",
        attempt=1,
        lease_id="lease-a",
        lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
        host_id="host-a",
        device_id="device-a",
        goal="open settings",
    )


def test_client_sends_authenticated_heartbeat_and_claim() -> None:
    requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        if request.url.path.endswith("/heartbeat"):
            return httpx.Response(
                200,
                json={
                    "host_id": "host-a",
                    "accepted_devices": 1,
                    "received_at": "2026-07-12T00:00:00Z",
                },
            )
        return httpx.Response(200, json={"assignment": None, "timed_out": True})

    async def scenario() -> None:
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://control.example",
        ) as http_client:
            client = HostAgentClient(_config(), http_client=http_client)
            heartbeat = await client.heartbeat(
                [
                    DeviceSnapshotModel(
                        device_id="device-a",
                        driver_type="wda",
                        status="idle",
                    )
                ]
            )
            assignment = await client.claim()
            assert heartbeat.accepted_devices == 1
            assert assignment is None

    asyncio.run(scenario())
    assert [request.url.path for request in requests] == [
        "/internal/v1/hosts/host-a/heartbeat",
        "/internal/v1/hosts/host-a/assignments/claim",
    ]
    assert all(
        request.headers["authorization"] == "Bearer host-secret" for request in requests
    )


def test_transport_failures_retry_with_bounded_backoff() -> None:
    attempts = 0
    sleeps: list[float] = []

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            raise httpx.ConnectError("temporary", request=request)
        return httpx.Response(200, json={"assignment": None, "timed_out": True})

    async def sleep(delay: float) -> None:
        sleeps.append(delay)

    async def scenario() -> None:
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://control.example",
        ) as http_client:
            client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
            assert await client.claim() is None

    asyncio.run(scenario())
    assert attempts == 3
    assert sleeps == [0.01, 0.02]


def test_stale_lease_response_raises_typed_error_without_retry() -> None:
    attempts = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal attempts
        attempts += 1
        return httpx.Response(
            409,
            json={"code": "stale_lease", "detail": "lease expired"},
        )

    async def scenario() -> None:
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://control.example",
        ) as http_client:
            client = HostAgentClient(_config(), http_client=http_client)
            with pytest.raises(StaleLeaseError) as error:
                await client.renew(_assignment())
            assert "host-secret" not in str(error.value)

    asyncio.run(scenario())
    assert attempts == 1