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, TaskScheduler
from cloud.store import CloudStore
from runtime.tool_calling_client import ToolCallDecision, ToolCallUsage


def _build_client(
    tmp_path,
    *,
    planner_client_factory=None,
    planner_token_reservation_ceiling: int = 4096,
) -> 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,
            scheduler=TaskScheduler(pool, pool.store, CloudConfig(stale_after_seconds=60)),
            planner_client_factory=planner_client_factory,
            planner_token_reservation_ceiling=planner_token_reservation_ceiling,
        )
    )
    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_heartbeat_records_host_planner_transport(tmp_path) -> None:
    client, pool = _build_client(tmp_path)

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

    assert response.status_code == 200
    assert pool.store.get_host("host-a").planner_transport == "cloud"  # type: ignore[union-attr]


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_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    headers = {"Authorization": "Bearer token-a"}
    heartbeat = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers=headers,
        json=_heartbeat_payload("host-a", "device-a"),
    )
    assert heartbeat.status_code == 200
    assert heartbeat.json()["policy_revision"] == 0
    assert heartbeat.json()["policy"] is None

    created = client.post(
        "/internal/v1/hosts/host-a/tasks",
        headers=headers,
        json={"host_id": "host-a", "goal": "local work", "device_id": "device-a"},
    )
    assert created.status_code == 201, created.text
    task = pool.store.get_task(created.json()["task_id"])
    assert task is not None
    assert task.constraints.target_host_id == "host-a"
    assert task.constraints.target_device_id == "device-a"
    foreign = client.post(
        "/internal/v1/hosts/host-b/tasks",
        headers=headers,
        json={"host_id": "host-b", "goal": "forbidden"},
    )
    assert foreign.status_code == 403


def test_host_policy_converges_and_disables_self_submission(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    headers = {"Authorization": "Bearer token-a"}
    client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers=headers,
        json=_heartbeat_payload("host-a", "device-a"),
    )
    pool.store.upsert_host_governance_policy(
        host_id="host-a",
        self_submission_enabled=False,
        max_active_tasks=2,
        daily_token_budget=1000,
        updated_at=datetime.now(UTC),
    )

    stale = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers=headers,
        json={**_heartbeat_payload("host-a", "device-a"), "policy_revision": 0},
    )
    assert stale.status_code == 200
    assert stale.json()["policy"] == {
        "revision": 1,
        "self_submission_enabled": False,
        "max_active_tasks": 2,
        "daily_token_budget": 1000,
    }
    current = client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers=headers,
        json={**_heartbeat_payload("host-a", "device-a"), "policy_revision": 1},
    )
    assert current.status_code == 200
    assert current.json()["policy"] is None
    disabled = client.post(
        "/internal/v1/hosts/host-a/tasks",
        headers=headers,
        json={"host_id": "host-a", "goal": "should be rejected"},
    )
    assert disabled.status_code == 403


def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -> None:
    class FakePlannerClient:
        calls = 0

        def decide(self, **_kwargs):
            self.calls += 1
            return ToolCallDecision(
                tool_name="tap",
                arguments={"x": 1, "y": 2},
                usage=ToolCallUsage(input_tokens=1, output_tokens=1, total_tokens=2),
            )

    planner = FakePlannerClient()
    client, pool = _build_client(
        tmp_path,
        planner_client_factory=lambda: planner,
        planner_token_reservation_ceiling=5,
    )
    headers = {"Authorization": "Bearer token-a"}
    client.put(
        "/internal/v1/hosts/host-a/heartbeat",
        headers=headers,
        json=_heartbeat_payload("host-a", "device-a"),
    )
    pool.store.upsert_host_governance_policy(
        host_id="host-a",
        self_submission_enabled=True,
        max_active_tasks=None,
        daily_token_budget=5,
        updated_at=datetime.now(UTC),
    )
    payload = {
        "host_id": "host-a",
        "system_prompt": "system",
        "user_prompt": "user",
        "tools": [{"name": "tap", "description": "tap", "parameters": {}}],
        "timeout_seconds": 1,
    }

    first = client.post(
        "/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
    )
    second = client.post(
        "/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
    )

    assert first.status_code == 200, first.text
    assert first.json()["total_tokens"] == 2
    assert second.status_code == 429
    assert planner.calls == 1
    events = pool.store.list_host_token_usage_events(
        host_id="host-a", limit=10, offset=0
    )
    assert len(events) == 1
    assert events[0].total_tokens == 2
    assert events[0].task_id 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"
    assert response.json()["cancel_requested"] is False
    renewed_expiry = datetime.fromisoformat(response.json()["lease_expires_at"])
    assert renewed_expiry > original_time + timedelta(seconds=30)


def test_lease_renewal_surfaces_pending_cancellation(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    _seed_active_assignment(pool)
    pool.store.request_task_cancellation("active-task", requested_at=datetime.now(UTC))

    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"
    assert response.json()["cancel_requested"] is True


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_cancelled_terminal_result_is_accepted_and_idempotent(tmp_path) -> None:
    client, pool = _build_client(tmp_path)
    _seed_active_assignment(pool)
    pool.store.request_task_cancellation("active-task", requested_at=datetime.now(UTC))
    payload = {
        "host_id": "host-a",
        "task_id": "active-task",
        "attempt": 1,
        "lease_id": "active-lease",
        "status": "cancelled",
        "failure_reason": "cancellation requested by control plane",
    }

    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 == "cancelled"  # 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