from datetime import datetime
from enum import StrEnum

from pydantic import BaseModel, model_validator


class Mount(BaseModel):
    src: str
    dst: str
    ro: bool = False


class ContainerState(StrEnum):
    RUNNING = "running"
    EXITED = "exited"
    DEAD = "dead"
    CREATED = "created"
    PAUSED = "paused"
    RESTARTING = "restarting"
    REMOVING = "removing"


class RunSpec(BaseModel):
    name: str
    image: str

    command: list[str] | None = None
    shell: str | None = None

    labels: dict[str, str] = {}
    gpus: list[int] | None = None
    cpu_cores: int | None = None
    cpuset: str | None = None
    memory_gb: int | None = None

    scratch: str | None = None
    mounts: list[Mount] = []

    network: str | None = None
    ports: dict[str, int] = {}
    workdir: str | None = None
    env: dict[str, str] = {}

    cap_drop_all: bool = True
    no_new_privileges: bool = True

    @model_validator(mode="after")
    def _check_command_xor_shell(self) -> "RunSpec":
        if self.command is not None and self.shell is not None:
            raise ValueError("command and shell are mutually exclusive")
        return self

    @model_validator(mode="after")
    def _check_cpu_cores_xor_cpuset(self) -> "RunSpec":
        if self.cpu_cores is not None and self.cpuset is not None:
            raise ValueError("cpu_cores and cpuset are mutually exclusive")
        return self


class RunResult(BaseModel):
    container_id: str
    name: str
    started_at: datetime


class ContainerStatus(BaseModel):
    container_id: str
    name: str
    state: ContainerState
    exit_code: int | None = None
    started_at: datetime | None = None
    finished_at: datetime | None = None
