"""Schema for NewJob, used to create new jobs."""

from msgspec import Struct, field

from slices_clientlib_aiis.schema.job import (
    DurationStr,
    JobRequestBase,
    JobRequestExtra,
    JobRequestScheduling,
    JobState,
    RestartInfo,
)
from slices_clientlib_aiis.util.msgspec import NonEmptyStr


class NewJobRequestScheduling(JobRequestScheduling):
    """GPULab scheduling instructions for new Jobs.

    This has extra defaults, and allowes DurationStr in not_after and not_before.
    """

    min_duration: DurationStr = "2 hour"
    max_duration: DurationStr = "8 hour"
    # This is what we want
    # not_before: AwareDatetime | DurationStr | None = None
    # not_after: AwareDatetime | DurationStr | None = "1 week"
    # But msgspec does not support unions with multiple str types, so we use plain str
    not_before: str | None = None  # pyright: ignore [reportIncompatibleVariableOverride] should be OK for this usage
    not_after: str | None = "1 week"  # pyright: ignore [reportIncompatibleVariableOverride] should be OK for this usage


class NewJobRequest(JobRequestBase):
    """Job Request for new jobs (extra defaults)."""

    scheduling: NewJobRequestScheduling = field(default_factory=lambda: NewJobRequestScheduling())
    extra: JobRequestExtra = field(default_factory=lambda: JobRequestExtra())


class NewJob(Struct, rename="camel"):
    """New job: schema limited to the values that must be present during initial submission."""

    name: NonEmptyStr
    request: NewJobRequest
    description: NonEmptyStr | None = None
    state: JobState | None = None
    restart_info: RestartInfo | None = None


class NewJobReq(Struct):
    """Request schema for new job creation."""

    job: NewJob
    idtoken: str
