from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, Literal, Protocol

if TYPE_CHECKING:
    from cloud.plugins import PluginManifest
    from cloud.pool import HostRegistration, PooledDevice
    from cloud.scheduler import ScheduledTask


AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
TerminalTaskStatus = Literal["done", "failed"]
ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
LeaseRenewalStatus = Literal["renewed", "not_found", "conflict", "expired"]


@dataclass(frozen=True)
class TaskAttemptRecord:
    task_id: str
    attempt: int
    lease_id: str
    host_id: str
    device_id: str
    status: AttemptStatus
    lease_expires_at: datetime
    created_at: datetime
    completed_at: datetime | None = None
    failure_reason: str | None = None
    terminal_result: dict[str, Any] | None = None


@dataclass(frozen=True)
class LeasedAssignment:
    task_id: str
    attempt: int
    lease_id: str
    lease_expires_at: datetime
    host_id: str
    device_id: str
    goal: str | None
    workflow_definition_id: str | None


class CloudRepository(Protocol):
    """Persistence port for cloud state and atomic scheduling operations."""

    def upsert_host(
        self,
        host_id: str,
        *,
        address: str | None,
        last_seen_at: datetime,
    ) -> None: ...

    def replace_host_devices(
        self,
        host_id: str,
        devices: list[PooledDevice],
    ) -> None: ...

    def list_hosts(self) -> list[HostRegistration]: ...

    def get_host(self, host_id: str) -> HostRegistration | None: ...

    def list_devices(self) -> list[PooledDevice]: ...

    def get_device(self, device_id: str) -> PooledDevice | None: ...

    def enqueue_task(self, task: ScheduledTask) -> None: ...

    def list_queued_tasks(self) -> list[ScheduledTask]: ...

    def get_task(self, task_id: str) -> ScheduledTask | None: ...

    def update_task(
        self,
        task_id: str,
        *,
        status: str | None = None,
        assigned_device_id: str | None = None,
        assigned_host_id: str | None = None,
    ) -> None: ...

    def count_queued_tasks(self) -> int: ...

    def save_plugin(self, manifest: PluginManifest, *, wired: bool) -> None: ...

    def list_plugins(self) -> list[tuple[PluginManifest, bool]]: ...

    def get_plugin(self, name: str) -> tuple[PluginManifest, bool] | None: ...

    def list_reserved_device_ids(self, *, now: datetime) -> set[str]: ...

    def assign_task(
        self,
        *,
        task_id: str,
        host_id: str,
        device_id: str,
        lease_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> LeasedAssignment | None: ...

    def claim_assignment(
        self,
        *,
        host_id: str,
        now: datetime,
    ) -> LeasedAssignment | None: ...

    def renew_lease(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> LeaseRenewalStatus: ...

    def record_task_result(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        status: TerminalTaskStatus,
        failure_reason: str | None,
        terminal_result: dict[str, Any] | None,
        completed_at: datetime,
    ) -> ResultRecordStatus: ...

    def reap_expired_leases(
        self,
        *,
        now: datetime,
        max_attempts: int,
    ) -> list[str]: ...

    def list_task_attempts(self, task_id: str) -> list[TaskAttemptRecord]: ...

    def health_check(self) -> None: ...

    def close(self) -> None: ...