from __future__ import annotations

from datetime import UTC, datetime

from fastapi import FastAPI
from fastapi.testclient import TestClient

from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider
from cloud.config import CloudConfig
from cloud.internal_api.api import create_internal_router
from cloud.pool import DevicePool, PooledDevice
from cloud.store import CloudStore


def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
    pool = DevicePool(
        CloudStore(tmp_path / "internal.sqlite3"),
        CloudConfig(stale_after_seconds=60),
    )
    auth_provider = ConfiguredBearerAuthProvider(
        [
            BearerCredential(
                principal_id="agent-a",
                token="token-a",
                host_id="host-a",
            ),
            BearerCredential(
                principal_id="agent-b",
                token="token-b",
                host_id="host-b",
            ),
        ]
    )
    app = FastAPI()
    app.include_router(create_internal_router(pool=pool, auth_provider=auth_provider))
    return TestClient(app), pool


def _heartbeat_payload(host_id: str, *device_ids: str) -> dict[str, object]:
    return {
        "host_id": host_id,
        "address": f"{host_id}.internal",
        "devices": [
            {
                "device_id": device_id,
                "driver_type": "wda",
                "status": "idle",
                "capability_tags": ["ios"],
            }
            for device_id in device_ids
        ],
    }


def test_authenticated_heartbeat_replaces_complete_snapshot(tmp_path) -> None:
    client, pool = _build_client(tmp_path)

    first = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers={"Authorization": "Bearer token-a"},
        json=_heartbeat_payload("host-a", "device-1", "device-2"),
    )
    second = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers={"Authorization": "Bearer token-a"},
        json=_heartbeat_payload("host-a", "device-2", "device-3"),
    )

    assert first.status_code == 200
    assert second.status_code == 200
    assert second.json()["accepted_devices"] == 2
    assert {device.device_id for device in pool.store.list_devices()} == {
        "device-2",
        "device-3",
    }


def test_invalid_duplicate_snapshot_preserves_previous_devices(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers={"Authorization": "Bearer token-a"},
        json=_heartbeat_payload("host-a", "existing-device"),
    )

    response = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers={"Authorization": "Bearer token-a"},
        json=_heartbeat_payload("host-a", "duplicate", "duplicate"),
    )

    assert response.status_code == 422
    assert [device.device_id for device in pool.store.list_devices()] == [
        "existing-device"
    ]


def test_live_device_owner_conflict_is_rejected_without_partial_sync(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    now = datetime.now(UTC)
    pool.store.upsert_host("host-a", address=None, last_seen_at=now)
    pool.store.replace_host_devices(
        "host-a",
        [
            PooledDevice(
                device_id="shared-device",
                host_id="host-a",
                driver_type="wda",
                status="idle",
                synced_at=now,
            )
        ],
    )

    response = client.put(
        "/internal/v1/hosts/host-b/heartbeat",
        headers={"Authorization": "Bearer token-b"},
        json=_heartbeat_payload("host-b", "shared-device"),
    )

    assert response.status_code == 409
    assert pool.store.get_host("host-b") is None
    devices = pool.store.list_devices()
    assert len(devices) == 1
    assert devices[0].host_id == "host-a"


def test_host_token_cannot_submit_heartbeat_for_another_host(tmp_path) -> None:
    client, pool = _build_client(tmp_path)

    response = client.put(
        "/internal/v1/hosts/host-b/heartbeat",
        headers={"Authorization": "Bearer token-a"},
        json=_heartbeat_payload("host-b", "device-b"),
    )

    assert response.status_code == 403
    assert pool.store.get_host("host-b") is None