from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, override

from ai.backend.common.types import DefaultForUnspecified, ResourceSlot
from ai.backend.manager.repositories.base.updater import UpdaterSpec
from ai.backend.manager.types import OptionalState, TriState

if TYPE_CHECKING:
    from ai.backend.manager.models.resource_policy import KeyPairResourcePolicyRow


@dataclass
class KeyPairResourcePolicyUpdaterSpec(UpdaterSpec["KeyPairResourcePolicyRow"]):
    """UpdaterSpec for keypair resource policy updates."""

    allowed_vfolder_hosts: OptionalState[dict[str, Any]] = field(default_factory=OptionalState.nop)
    default_for_unspecified: OptionalState[DefaultForUnspecified] = field(
        default_factory=OptionalState.nop
    )
    idle_timeout: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_concurrent_sessions: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_containers_per_session: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_pending_session_count: TriState[int] = field(default_factory=TriState.nop)
    max_pending_session_resource_slots: TriState[dict[str, Any]] = field(
        default_factory=TriState.nop
    )
    max_quota_scope_size: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_vfolder_count: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_vfolder_size: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_concurrent_sftp_sessions: OptionalState[int] = field(default_factory=OptionalState.nop)
    max_session_lifetime: OptionalState[int] = field(default_factory=OptionalState.nop)
    total_resource_slots: OptionalState[ResourceSlot] = field(default_factory=OptionalState.nop)

    @property
    @override
    def row_class(self) -> type[KeyPairResourcePolicyRow]:
        from ai.backend.manager.models.resource_policy import KeyPairResourcePolicyRow

        return KeyPairResourcePolicyRow

    @override
    def build_values(self) -> dict[str, Any]:
        to_update: dict[str, Any] = {}
        self.allowed_vfolder_hosts.update_dict(to_update, "allowed_vfolder_hosts")
        self.default_for_unspecified.update_dict(to_update, "default_for_unspecified")
        self.idle_timeout.update_dict(to_update, "idle_timeout")
        self.max_concurrent_sessions.update_dict(to_update, "max_concurrent_sessions")
        self.max_containers_per_session.update_dict(to_update, "max_containers_per_session")
        self.max_pending_session_count.update_dict(to_update, "max_pending_session_count")
        self.max_pending_session_resource_slots.update_dict(
            to_update, "max_pending_session_resource_slots"
        )
        self.max_quota_scope_size.update_dict(to_update, "max_quota_scope_size")
        self.max_vfolder_count.update_dict(to_update, "max_vfolder_count")
        self.max_vfolder_size.update_dict(to_update, "max_vfolder_size")
        self.max_concurrent_sftp_sessions.update_dict(to_update, "max_concurrent_sftp_sessions")
        self.max_session_lifetime.update_dict(to_update, "max_session_lifetime")
        self.total_resource_slots.update_dict(to_update, "total_resource_slots")
        return to_update
