from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta
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, ScheduledTaskStatus
    from cloud.user_auth import (
        AuthAuditEvent,
        AuthenticatedUserSession,
        LoginThrottle,
        UserAccount,
        UserSession,
    )
    from cloud.governance import HostGovernancePolicy, UserSubmissionPolicy


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


class HostEnrollmentConflictError(RuntimeError):
    """Raised when a Host enrollment cannot preserve its existing binding."""


class EnrollmentTokenConflictError(HostEnrollmentConflictError):
    """Raised when a one-time enrollment token is already bound elsewhere."""


class DeviceEnrollmentConflictError(RuntimeError):
    """Raised when a local device enrollment conflicts with stored identity."""


class UserConflictError(RuntimeError):
    """Raised when a user operation violates a durable identity invariant."""


class LastAdministratorConflictError(UserConflictError):
    """Raised when a write would remove the last enabled administrator."""


@dataclass(frozen=True)
class HostEnrollment:
    host_id: str
    agent_instance_id: str
    display_name: str | None
    enrolled_at: datetime
    revoked_at: datetime | None = None


@dataclass(frozen=True)
class DeviceEnrollment:
    device_id: str
    host_id: str
    local_device_id: str
    driver_type: str
    name: str | None
    capability_tags: list[str]
    enrolled_at: datetime
    revoked_at: datetime | None = None


@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 enroll_host(
        self,
        *,
        host_id: str,
        agent_instance_id: str,
        credential_digest: str,
        enrollment_token_digest: str | None,
        display_name: str | None,
        enrolled_at: datetime,
    ) -> HostEnrollment: ...

    def authenticate_enrolled_host(
        self,
        credential_digest: str,
    ) -> str | None: ...

    def revoke_enrolled_host(
        self,
        host_id: str,
        *,
        revoked_at: datetime,
    ) -> bool: ...

    def is_enrollment_managed_host(self, host_id: str) -> bool: ...

    def enroll_device(
        self,
        *,
        device_id: str,
        host_id: str,
        local_device_id: str,
        driver_type: str,
        name: str | None,
        capability_tags: list[str],
        enrolled_at: datetime,
    ) -> DeviceEnrollment: ...

    def get_device_enrollment(
        self,
        device_id: str,
    ) -> DeviceEnrollment | None: ...

    def list_device_enrollments(
        self,
        host_id: str,
    ) -> list[DeviceEnrollment]: ...

    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],
        *,
        allow_device_takeover: bool = False,
    ) -> 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 list_tasks(
        self,
        *,
        status: ScheduledTaskStatus | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[ScheduledTask]: ...

    def count_tasks(self, status: ScheduledTaskStatus | None = None) -> int: ...

    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 create_user(self, user: UserAccount) -> UserAccount: ...

    def get_user(self, user_id: str) -> UserAccount | None: ...

    def get_user_by_normalized_username(
        self,
        username_normalized: str,
    ) -> UserAccount | None: ...

    def list_users(self, *, limit: int, offset: int) -> list[UserAccount]: ...

    def update_user(
        self,
        user_id: str,
        *,
        display_name: str | None = None,
        role: str | None = None,
        enabled: bool | None = None,
        updated_at: datetime,
    ) -> UserAccount: ...

    def rehash_user_password(
        self,
        user_id: str,
        *,
        password_hash: str,
        updated_at: datetime,
    ) -> UserAccount: ...

    def update_user_password(
        self,
        user_id: str,
        *,
        password_hash: str,
        must_change_password: bool,
        updated_at: datetime,
        revoke_sessions: bool,
    ) -> UserAccount: ...

    def mark_user_login(self, user_id: str, *, now: datetime) -> UserAccount: ...

    def create_user_session(self, session: UserSession) -> None: ...

    def get_authenticated_user_session(
        self,
        token_digest: str,
        *,
        now: datetime,
    ) -> AuthenticatedUserSession | None: ...

    def touch_user_session(
        self,
        session_id: str,
        *,
        last_seen_at: datetime,
        idle_expires_at: datetime,
    ) -> UserSession: ...

    def revoke_user_session(self, session_id: str, *, revoked_at: datetime) -> bool: ...

    def revoke_user_sessions(self, user_id: str, *, revoked_at: datetime) -> int: ...

    def get_login_throttle(
        self,
        username_normalized: str,
        client_bucket: str,
    ) -> LoginThrottle | None: ...

    def record_login_failure(
        self,
        *,
        username_normalized: str,
        client_bucket: str,
        now: datetime,
        failure_limit: int,
        failure_window: timedelta,
        block_duration: timedelta,
    ) -> LoginThrottle: ...

    def clear_login_throttle(self, username_normalized: str, client_bucket: str) -> None: ...

    def record_auth_audit(self, event: AuthAuditEvent) -> None: ...

    def cleanup_auth_state(self, *, now: datetime, limit: int) -> int: ...

    def get_user_submission_policy(self, user_id: str) -> UserSubmissionPolicy | None: ...

    def upsert_user_submission_policy(
        self,
        *,
        user_id: str,
        submission_enabled: bool,
        allowed_host_ids: tuple[str, ...] | None,
        allowed_device_targets: tuple[tuple[str, str], ...] | None,
        updated_at: datetime,
    ) -> UserSubmissionPolicy: ...

    def get_host_governance_policy(self, host_id: str) -> HostGovernancePolicy | None: ...

    def upsert_host_governance_policy(
        self,
        *,
        host_id: str,
        self_submission_enabled: bool,
        max_active_tasks: int | None,
        daily_token_budget: int | None,
        updated_at: datetime,
    ) -> HostGovernancePolicy: ...

    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: ...