from __future__ import annotations

from datetime import UTC, datetime
from datetime import timedelta

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.scheduler import ScheduledTask, TaskConstraints
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_empty_heartbeat_removes_only_reporting_hosts_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", "device-a"),
    )
    client.put(
        "/internal/v1/hosts/host-b/heartbeat",
        headers={"Authorization": "Bearer token-b"},
        json=_heartbeat_payload("host-b", "device-b"),
    )

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

    assert response.status_code == 200
    assert response.json()["accepted_devices"] == 0
    devices = pool.store.list_devices()
    assert [(device.device_id, device.host_id) for device in devices] == [
        ("device-b", "host-b")
    ]


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_stale_device_owner_can_be_replaced_by_live_host(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    stale_at = datetime.now(UTC) - timedelta(seconds=61)
    pool.store.upsert_host("host-a", address=None, last_seen_at=stale_at)
    pool.store.replace_host_devices(
        "host-a",
        [
            PooledDevice(
                device_id="shared-device",
                host_id="host-a",
                driver_type="wda",
                status="idle",
                synced_at=stale_at,
            )
        ],
    )

    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 == 200
    device = pool.store.list_devices()[0]
    assert device.device_id == "shared-device"
    assert device.host_id == "host-b"


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


def test_long_poll_claim_returns_at_most_one_owned_assignment(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="claim-device",
                host_id="host-a",
                driver_type="wda",
                status="idle",
                synced_at=now,
            )
        ],
    )
    pool.store.enqueue_task(
        ScheduledTask(
            id="claim-task",
            goal="open settings",
            workflow_definition_id=None,
            constraints=TaskConstraints(),
            created_at=now,
        )
    )
    pool.store.assign_task(
        task_id="claim-task",
        host_id="host-a",
        device_id="claim-device",
        lease_id="claim-lease",
        lease_expires_at=now + timedelta(minutes=1),
        now=now,
    )

    first = client.post(
        "/internal/v1/hosts/host-a/assignments/claim",
        headers={"Authorization": "Bearer token-a"},
        json={"host_id": "host-a", "timeout_seconds": 0},
    )
    second = client.post(
        "/internal/v1/hosts/host-a/assignments/claim",
        headers={"Authorization": "Bearer token-a"},
        json={"host_id": "host-a", "timeout_seconds": 0},
    )

    assert first.status_code == 200
    assignment = first.json()["assignment"]
    assert assignment["task_id"] == "claim-task"
    assert assignment["lease_id"] == "claim-lease"
    assert second.status_code == 200
    assert second.json() == {"assignment": None, "timed_out": True}


def test_empty_long_poll_timeout_is_normal_response(tmp_path) -> None:
    client, _ = _build_client(tmp_path)

    response = client.post(
        "/internal/v1/hosts/host-a/assignments/claim",
        headers={"Authorization": "Bearer token-a"},
        json={"host_id": "host-a", "timeout_seconds": 0},
    )

    assert response.status_code == 200
    assert response.json() == {"assignment": None, "timed_out": True}


def _seed_active_assignment(pool: DevicePool) -> datetime:
    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="active-device",
                host_id="host-a",
                driver_type="wda",
                status="idle",
                synced_at=now,
            )
        ],
    )
    pool.store.enqueue_task(
        ScheduledTask(
            id="active-task",
            goal="execute assignment",
            workflow_definition_id=None,
            constraints=TaskConstraints(),
            created_at=now,
        )
    )
    pool.store.assign_task(
        task_id="active-task",
        host_id="host-a",
        device_id="active-device",
        lease_id="active-lease",
        lease_expires_at=now + timedelta(minutes=1),
        now=now,
    )
    pool.store.claim_assignment(host_id="host-a", now=now)
    return now


def test_lease_renewal_extends_active_assignment(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    original_time = _seed_active_assignment(pool)

    response = client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/renew",
        headers={"Authorization": "Bearer token-a"},
        json={
            "host_id": "host-a",
            "task_id": "active-task",
            "attempt": 1,
            "lease_id": "active-lease",
        },
    )

    assert response.status_code == 200
    assert response.json()["status"] == "renewed"
    renewed_expiry = datetime.fromisoformat(response.json()["lease_expires_at"])
    assert renewed_expiry > original_time + timedelta(seconds=30)


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

    response = client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/renew",
        headers={"Authorization": "Bearer token-a"},
        json={
            "host_id": "host-a",
            "task_id": "active-task",
            "attempt": 1,
            "lease_id": "stale-lease",
        },
    )

    assert response.status_code == 409
    assert response.json()["code"] == "stale_lease"


def test_terminal_result_is_idempotent_through_internal_api(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    _seed_active_assignment(pool)
    payload = {
        "host_id": "host-a",
        "task_id": "active-task",
        "attempt": 1,
        "lease_id": "active-lease",
        "status": "done",
        "result": {"steps": 4},
    }

    first = client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/result",
        headers={"Authorization": "Bearer token-a"},
        json=payload,
    )
    repeated = client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/result",
        headers={"Authorization": "Bearer token-a"},
        json=payload,
    )

    assert first.status_code == 200
    assert first.json()["status"] == "recorded"
    assert repeated.status_code == 200
    assert repeated.json()["status"] == "already_recorded"
    assert pool.store.get_task("active-task").status == "done"  # type: ignore[union-attr]


def test_conflicting_repeated_result_returns_stale_lease_conflict(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    _seed_active_assignment(pool)
    base_payload = {
        "host_id": "host-a",
        "task_id": "active-task",
        "attempt": 1,
        "lease_id": "active-lease",
        "status": "done",
        "result": {"steps": 4},
    }
    client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/result",
        headers={"Authorization": "Bearer token-a"},
        json=base_payload,
    )

    response = client.post(
        "/internal/v1/hosts/host-a/assignments/active-task/result",
        headers={"Authorization": "Bearer token-a"},
        json={**base_payload, "status": "failed", "failure_reason": "late failure"},
    )

    assert response.status_code == 409
    assert response.json()["code"] == "stale_lease"


def test_multi_host_claims_are_isolated(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    now = datetime.now(UTC)
    for host_id, device_id, task_id, lease_id in (
        ("host-a", "device-a", "task-a", "lease-a"),
        ("host-b", "device-b", "task-b", "lease-b"),
    ):
        pool.store.upsert_host(host_id, address=None, last_seen_at=now)
        pool.store.replace_host_devices(
            host_id,
            [
                PooledDevice(
                    device_id=device_id,
                    host_id=host_id,
                    driver_type="wda",
                    status="idle",
                    synced_at=now,
                )
            ],
        )
        pool.store.enqueue_task(
            ScheduledTask(
                id=task_id,
                goal=f"work for {host_id}",
                workflow_definition_id=None,
                constraints=TaskConstraints(),
                created_at=now,
            )
        )
        pool.store.assign_task(
            task_id=task_id,
            host_id=host_id,
            device_id=device_id,
            lease_id=lease_id,
            lease_expires_at=now + timedelta(minutes=1),
            now=now,
        )

    host_a = client.post(
        "/internal/v1/hosts/host-a/assignments/claim",
        headers={"Authorization": "Bearer token-a"},
        json={"host_id": "host-a", "timeout_seconds": 0},
    )
    host_b = client.post(
        "/internal/v1/hosts/host-b/assignments/claim",
        headers={"Authorization": "Bearer token-b"},
        json={"host_id": "host-b", "timeout_seconds": 0},
    )

    assert host_a.json()["assignment"]["task_id"] == "task-a"
    assert host_b.json()["assignment"]["task_id"] == "task-b"


def test_superseded_attempt_result_cannot_overwrite_current_lease(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="retry-device",
                host_id="host-a",
                driver_type="wda",
                status="idle",
                synced_at=now,
            )
        ],
    )
    pool.store.enqueue_task(
        ScheduledTask(
            id="retry-task",
            goal="retry safely",
            workflow_definition_id=None,
            constraints=TaskConstraints(),
            created_at=now,
        )
    )
    pool.store.assign_task(
        task_id="retry-task",
        host_id="host-a",
        device_id="retry-device",
        lease_id="old-lease",
        lease_expires_at=now + timedelta(seconds=1),
        now=now,
    )
    reaped_at = now + timedelta(seconds=2)
    pool.store.reap_expired_leases(now=reaped_at, max_attempts=2)
    pool.store.assign_task(
        task_id="retry-task",
        host_id="host-a",
        device_id="retry-device",
        lease_id="current-lease",
        lease_expires_at=reaped_at + timedelta(minutes=1),
        now=reaped_at,
    )

    response = client.post(
        "/internal/v1/hosts/host-a/assignments/retry-task/result",
        headers={"Authorization": "Bearer token-a"},
        json={
            "host_id": "host-a",
            "task_id": "retry-task",
            "attempt": 1,
            "lease_id": "old-lease",
            "status": "done",
            "result": {"late": True},
        },
    )

    assert response.status_code == 409
    assert response.json()["code"] == "stale_lease"
    task = pool.store.get_task("retry-task")
    assert task is not None
    assert task.status == "assigned"
    assert task.attempt_count == 2
    assert task.lease_id == "current-lease"
    assert task.terminal_result is None