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
from host_agent.progress import TaskProgressSnapshot

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
        self._host_policy: dict[str, Any] | None = None
        self._latest_progress: TaskProgressSnapshot | 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
            self._latest_progress = None

    def set_latest_progress(self, snapshot: TaskProgressSnapshot | None) -> None:
        with self._lock:
            self._latest_progress = snapshot

    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 mark_host_policy(self, policy: Any | None) -> None:
        with self._lock:
            self._host_policy = (
                {
                    "revision": policy.revision,
                    "self_submission_enabled": policy.self_submission_enabled,
                    "max_active_tasks": policy.max_active_tasks,
                    "daily_token_budget": policy.daily_token_budget,
                }
                if policy is not None
                else None
            )

    def snapshot(self) -> dict[str, Any]:
        with self._lock:
            current_assignment = self._current_assignment
            last_heartbeat = self._last_heartbeat
            progress = self._latest_progress
        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
            ),
            "host_policy": self._host_policy.copy() if self._host_policy else None,
            "progress": (
                {
                    "step_index": progress.step_index,
                    "step_status": progress.step_status,
                    "summary": progress.summary,
                    "updated_at": progress.updated_at.isoformat(),
                }
                if progress is not None
                else None
            ),
        }