from __future__ import annotations

from datetime import UTC, datetime

import pytest
from pydantic import ValidationError

from cloud.internal_api.models import (
    AssignmentModel,
    ClaimRequest,
    ClaimResponse,
    DeviceSnapshotModel,
    HeartbeatRequest,
    LeaseRenewalRequest,
    TerminalResultRequest,
)


def test_heartbeat_models_complete_device_snapshot() -> None:
    heartbeat = HeartbeatRequest(
        host_id="host-a",
        address="10.0.0.1:9000",
        devices=[
            DeviceSnapshotModel(
                device_id="device-a",
                driver_type="wda",
                status="idle",
                capability_tags=["ios", "physical"],
            )
        ],
    )

    assert heartbeat.host_id == "host-a"
    assert heartbeat.devices[0].capability_tags == ["ios", "physical"]


@pytest.mark.parametrize("status", ["unreachable", "unknown", ""])
def test_device_snapshot_rejects_invalid_host_status(status: str) -> None:
    with pytest.raises(ValidationError):
        DeviceSnapshotModel(
            device_id="device-a",
            driver_type="wda",
            status=status,
        )


def test_claim_response_supports_assignment_and_normal_timeout() -> None:
    assignment = 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",
    )

    assert ClaimResponse(assignment=assignment).assignment == assignment
    assert ClaimResponse(timed_out=True).model_dump() == {
        "assignment": None,
        "timed_out": True,
    }


@pytest.mark.parametrize("timeout_seconds", [-1, 61])
def test_claim_timeout_is_bounded(timeout_seconds: float) -> None:
    with pytest.raises(ValidationError):
        ClaimRequest(host_id="host-a", timeout_seconds=timeout_seconds)


def test_renewal_and_result_require_attempt_and_lease_identity() -> None:
    renewal = LeaseRenewalRequest(
        host_id="host-a",
        task_id="task-a",
        attempt=2,
        lease_id="lease-a",
    )
    result = TerminalResultRequest(
        host_id="host-a",
        task_id="task-a",
        attempt=2,
        lease_id="lease-a",
        status="failed",
        failure_reason="device offline",
        result={"step": 3},
    )

    assert renewal.attempt == result.attempt == 2
    assert result.failure_reason == "device offline"


@pytest.mark.parametrize(
    "payload",
    [
        {"host_id": "", "task_id": "task-a", "attempt": 1, "lease_id": "lease"},
        {"host_id": "host-a", "task_id": "", "attempt": 1, "lease_id": "lease"},
        {"host_id": "host-a", "task_id": "task-a", "attempt": 0, "lease_id": "lease"},
        {"host_id": "host-a", "task_id": "task-a", "attempt": 1, "lease_id": ""},
    ],
)
def test_renewal_rejects_incomplete_identity(payload: dict[str, object]) -> None:
    with pytest.raises(ValidationError):
        LeaseRenewalRequest.model_validate(payload)