from __future__ import annotations

import threading
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from cloud.internal_api.models import AssignmentModel

if TYPE_CHECKING:
    from collections.abc import Callable


@dataclass(frozen=True)
class _CurrentAssignment:
    task_id: str
    device_id: str
    goal: str | None
    workflow_definition_id: str | None
    started_at: datetime


@dataclass(frozen=True)
class _LastHeartbeat:
    ok: bool
    device_count: int
    at: datetime


class AgentStatusTracker:
    def __init__(self, *, now: Callable[[], datetime] | None = None) -> None:
        self._now = now or (lambda: datetime.now(UTC))
        self._lock = threading.Lock()
        self._current_assignment: _CurrentAssignment | None = None
        self._last_heartbeat: _LastHeartbeat | None = None

    def mark_assignment_started(self, assignment: AssignmentModel) -> None:
        with self._lock:
            self._current_assignment = _CurrentAssignment(
                task_id=assignment.task_id,
                device_id=assignment.device_id,
                goal=assignment.goal,
                workflow_definition_id=assignment.workflow_definition_id,
                started_at=self._now(),
            )

    def mark_assignment_finished(self) -> None:
        with self._lock:
            self._current_assignment = None

    def mark_heartbeat(self, *, ok: bool, device_count: int) -> None:
        with self._lock:
            self._last_heartbeat = _LastHeartbeat(
                ok=ok,
                device_count=device_count,
                at=self._now(),
            )

    def snapshot(self) -> dict[str, Any]:
        with self._lock:
            current_assignment = self._current_assignment
            last_heartbeat = self._last_heartbeat
        return {
            "current_assignment": (
                {
                    "task_id": current_assignment.task_id,
                    "device_id": current_assignment.device_id,
                    "goal": current_assignment.goal,
                    "workflow_definition_id": current_assignment.workflow_definition_id,
                    "started_at": current_assignment.started_at.isoformat(),
                }
                if current_assignment is not None
                else None
            ),
            "last_heartbeat": (
                {
                    "ok": last_heartbeat.ok,
                    "device_count": last_heartbeat.device_count,
                    "at": last_heartbeat.at.isoformat(),
                }
                if last_heartbeat is not None
                else None
            ),
        }