from __future__ import annotations

from datetime import UTC, datetime

from cloud.internal_api.models import AssignmentModel
from host_agent.status import AgentStatusTracker


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",
        workflow_definition_id="workflow-a",
    )


def test_mark_assignment_started_reflected_in_snapshot() -> None:
    ticks = iter([datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC)])
    tracker = AgentStatusTracker(now=lambda: next(ticks))

    tracker.mark_assignment_started(_assignment())

    snapshot = tracker.snapshot()
    assert snapshot["current_assignment"] == {
        "task_id": "task-a",
        "device_id": "device-a",
        "goal": "open settings",
        "workflow_definition_id": "workflow-a",
        "started_at": "2026-07-13T09:00:00+00:00",
    }


def test_mark_assignment_finished_clears_current_assignment() -> None:
    tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))

    tracker.mark_assignment_started(_assignment())
    tracker.mark_assignment_finished()

    assert tracker.snapshot()["current_assignment"] is None


def test_mark_heartbeat_reflected_in_snapshot() -> None:
    tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 5, 0, tzinfo=UTC))

    tracker.mark_heartbeat(ok=True, device_count=3)

    snapshot = tracker.snapshot()
    assert snapshot["last_heartbeat"] == {
        "ok": True,
        "device_count": 3,
        "at": "2026-07-13T09:05:00+00:00",
    }


def test_snapshot_defaults_to_no_assignment_or_heartbeat() -> None:
    tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))

    snapshot = tracker.snapshot()

    assert snapshot["current_assignment"] is None
    assert snapshot["last_heartbeat"] is None