"""Schema for Worker/Slave and Cluster info."""

from msgspec import Struct, field

from slices_clientlib_aiis.util.msgspec import AwareDatetime


class GpuModel(Struct, rename="camel"):
    """GPU Model."""

    vendor: str
    name: str
    memory_mb: int


class ResourceInfo(Struct, rename="camel"):
    """Resource Information."""

    system_total: int
    acquired: int
    used: int
    available: int

    @classmethod
    def zero(cls):
        """Return instance with everything set to 0."""
        return cls(
            system_total=0,
            acquired=0,
            used=0,
            available=0,
        )


class ActiveJobsCounts(Struct, rename="camel"):
    """Return the number of jobs in any of the 'active' job states.

    Can be used for both clusters and slaves.
    """

    onhold: int
    queued: int
    assigned: int
    starting: int
    running: int
    musthalt: int
    halting: int

    @classmethod
    def zero(cls):
        """Return an ActiveJobsCounts with all counters on zero."""
        return cls(
            onhold=0,
            queued=0,
            assigned=0,
            starting=0,
            running=0,
            musthalt=0,
            halting=0,
        )


class Slave(Struct, rename="camel"):
    """Slave/Worker Info."""

    name: str
    instance_id: str
    cluster_id: int
    pid: int | None
    software_version: str | None

    aliases: list[str]
    comment: str | None
    host: str | None

    gpu_model: list[GpuModel]
    cpu_model: list[str]
    worker: ResourceInfo
    cpu_memory_mb: ResourceInfo
    gpu: ResourceInfo
    cpu: ResourceInfo
    last_update: AwareDatetime
    shutting_down: bool
    cuda_version_full: str | None = None
    cuda_version_major: int | None = None

    docker_disk_used_percent: float = -1
    accepting_jobs: bool = True
    active_jobs_counts: ActiveJobsCounts | None = None

    storage_paths_available: list[str] | None = None
    storage_aliases_available: dict[str, str] | None = None
    storage_paths_available_project_limited: dict[str, list[str]] | None = None
    active_job_uuids: list[str] | None = None


class Cluster(Struct, rename="camel"):
    """Cluster Info."""

    cluster_id: int
    site_id: str
    comment: str
    is_private: bool
    slave_count: int = 0
    gpu_model: list[GpuModel] = []
    worker: ResourceInfo = field(default_factory=ResourceInfo.zero)
    cpu_memory_mb: ResourceInfo = field(default_factory=ResourceInfo.zero)
    gpu: ResourceInfo = field(default_factory=ResourceInfo.zero)
    cpu: ResourceInfo = field(default_factory=ResourceInfo.zero)
    active_jobs_counts: ActiveJobsCounts | None = None
    have_access: bool | None = None
    allowed_projects: list[str] | None = None
    storage_paths_available: list[str] | None = None
    storage_aliases_available: dict | None = None
