"""Schema for Job."""

from datetime import datetime, timedelta
from enum import Enum
from typing import Annotated, Optional
from uuid import UUID

import pytz
from msgspec import Meta, Struct, field

from slices_clientlib_aiis.schema.job_statistics import GPULabUsageStatistics, WasteReview
from slices_clientlib_aiis.util.fabric import connect as connect_with_fabric
from slices_clientlib_aiis.util.msgspec import AwareDatetime, EmailStr, NonEmptyStr, PositiveInt, PublicSshKey

DurationStr = Annotated[str, Meta(min_length=2)]


class JobStatus(Enum):
    """Job state."""

    ONHOLD = "On Hold"
    QUEUED = "Queued"
    ASSIGNED = "Assigned"
    STARTING = "Starting"
    RUNNING = "Running"
    CANCELLING = "Cancelling"
    CANCELLED = "Cancelled"
    FINISHED = "Finished"
    DELETED = "Deleted"
    FAILED = "Failed"
    MUSTHALT = "MustHalt"
    HALTING = "Halting"
    HALTED = "Halted"

    @classmethod
    def _missing_(cls, value):
        """Handle missing enum values by ignoring case."""
        if isinstance(value, cls):
            return value
        if not isinstance(value, str):
            return None
        for member in cls:
            if member.name.lower() == value.lower():
                return member
            if member.value.lower() == value.lower():
                return member
        return None


class JobStorage(Struct, rename="camel"):
    """Job Storage (tmpfs or dictionary)."""

    container_path: NonEmptyStr | None = None
    host_path: NonEmptyStr | None = None
    size_gb: PositiveInt | None = None  # only for tmpfs


class JobPortMapping(Struct, rename="camel"):
    """Job Port Mapping."""

    container_port: int
    host_port: int | None = None
    host_ip: str | None = None


class JobRequestDocker(Struct, rename="camel"):
    """Job Request: Docker."""

    image: NonEmptyStr
    command: str | list[str] = []
    environment: dict[NonEmptyStr, str | int | bool] = {}
    storage: list[JobStorage] = []
    port_mappings: list[JobPortMapping] = []
    project_gid_variable_name: NonEmptyStr | None = None
    working_dir: NonEmptyStr | None = None
    group_add: list[NonEmptyStr] = []
    user: NonEmptyStr | None = None


class MaxSimultaneousJobs(Struct, rename="camel"):
    """Maximum number of simultaneous jobs."""

    bucket_name: str
    bucket_max: int


class JobRequestScheduling(Struct, rename="camel"):
    """GPULab scheduling instructions Base."""

    min_duration: DurationStr
    max_duration: DurationStr
    interactive: bool = False
    restartable: bool = False
    allow_halt_without_signal: bool = False
    reservation_ids: list[UUID] = []
    max_simultaneous_jobs: MaxSimultaneousJobs | None = None
    not_before: AwareDatetime | None = None
    not_after: AwareDatetime | None = None


class JobRequestResources(Struct, rename="camel"):
    """Job Request Resources."""

    cpus: PositiveInt
    gpus: int
    cpu_memory_gb: PositiveInt
    gpu_memory_gb: PositiveInt | None = None

    features: list[NonEmptyStr] = []
    gpu_model: list[NonEmptyStr] = []

    min_cuda_version: PositiveInt | None = None
    cluster_id: PositiveInt | None = None
    slave_name: NonEmptyStr | None = None
    slave_instance_id: NonEmptyStr | None = None


class JobRequestExtra(Struct, rename="camel"):
    """Job Request Extras."""

    ssh_pub_keys: list[PublicSshKey] = []
    email_on_queue: list[EmailStr] = []
    email_on_run: list[EmailStr] = []
    email_on_end: list[EmailStr] = []
    email_on_halt: list[EmailStr] = []
    email_on_restart: list[EmailStr] = []


class JobRequestBase(Struct, rename="camel"):
    """Job Request Base."""

    resources: JobRequestResources
    docker: JobRequestDocker


class JobRequest(JobRequestBase):
    """Job Request Base."""

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


class JobOwner(Struct, rename="camel"):
    """Job Owner."""

    user_id: NonEmptyStr
    project_id: NonEmptyStr
    experiment_id: NonEmptyStr | None


class RestartInfo(Struct, rename="camel"):
    """Restart Information.

    Available in jobs that have been restarted after being halted.
    """

    initial_job_uuid: UUID
    restart_count: int


class HaltReason(Enum):
    """Halt Reason."""

    UNKNOWN = "Unknown"
    USER = "User requested fallback manually"
    ADMIN = "Admin requested fallback"
    SCORE = "Score"
    RESERVATION = "Reservation Priority"
    CLUSTER = "Private Cluster Priority"

    @classmethod
    def _missing_(cls, value):
        """Handle missing enum values by ignoring case."""
        if isinstance(value, cls):
            return value
        if not isinstance(value, str):
            return None
        for member in cls:
            if member.name.lower() == value.lower():
                return member
            if member.value.lower() == value.lower():
                return member
        return None


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

    reason: HaltReason
    time: datetime
    started_job_uuids: list[str]
    halted_job_uuids: list[str]
    slave_name: str
    slave_instance: str
    started_job_scores: list[float]
    halted_job_scores: list[float]


class JobStateScheduling(Struct, rename="camel"):
    """Job State Scheduling."""

    assigned_cluster_id: int | None = None
    assigned_instance_id: str | None = None
    assigned_slave_name: str | None = None
    queued_explanations: list[str] = []
    tally_increment: float | None = None
    halt_events: list[HaltInfo] = []
    ignoring_reservation_ids: list[str] = []


class JobEventTimes(Struct, rename="camel"):
    """Job Event Times."""

    created: AwareDatetime
    status_updated: AwareDatetime
    QUEUED: AwareDatetime | None = field(name="QUEUED", default=None)
    ASSIGNED: AwareDatetime | None = field(name="ASSIGNED", default=None)
    STARTING: AwareDatetime | None = field(name="STARTING", default=None)
    RUNNING: AwareDatetime | None = field(name="RUNNING", default=None)
    FINISHED: AwareDatetime | None = field(name="FINISHED", default=None)
    FAILED: AwareDatetime | None = field(name="FAILED", default=None)
    CANCELLED: AwareDatetime | None = field(name="CANCELLED", default=None)
    DELETED: AwareDatetime | None = field(name="DELETED", default=None)
    MUSTHALT: AwareDatetime | None = field(name="MUSTHALT", default=None)
    HALTING: AwareDatetime | None = field(name="HALTING", default=None)
    HALTED: AwareDatetime | None = field(name="HALTED", default=None)
    long_run_notify: AwareDatetime | None = None
    waste_notify: AwareDatetime | None = None

    @property
    def end_date(self) -> AwareDatetime | None:
        """End date."""
        if self.FINISHED:
            return self.FINISHED
        if self.HALTED:
            return self.HALTED
        if self.FAILED:
            return self.FAILED
        if self.CANCELLED:
            return self.CANCELLED
        if self.DELETED:
            return self.DELETED
        return None

    @property
    def duration_or_none(self) -> timedelta | None:
        """Get current duration of job, or None if not started."""
        start = self.RUNNING
        if start is None:
            return None

        end = self.FINISHED or self.CANCELLED or self.FAILED or self.HALTED or self.DELETED
        if end is None:
            end = datetime.now(pytz.utc)
        return end - start

    @property
    def duration_or_zero(self) -> timedelta:
        """Get current duration of job, or 0 duration if not started."""
        return self.duration_or_none or timedelta(seconds=0)

    @property
    def queue_duration(self) -> timedelta:
        """Get queue duration, or zero if not ever queued."""
        start = self.QUEUED
        if start is None:
            return timedelta(seconds=0)

        end = self.STARTING or self.RUNNING or self.FAILED or self.HALTED or self.CANCELLED or self.DELETED
        if end is None:
            end = datetime.now(pytz.utc)

        return end - start


class GpuInfo(Struct, rename="camel"):
    """GPU Information."""

    index: int
    uuid: str
    serial: str
    name: str
    brand: str
    minor_number: int
    board_id: int
    bridge_chip_info: str
    is_multi_gpu_board: bool
    max_pcie_link_generation: int
    max_pcie_link_width: int
    vbios_version: str


class GpuOverview(Struct, rename="camel"):
    """GPU Information Overview."""

    cuda_version_full: str
    cuda_version_int: int
    cuda_version_major: int
    cuda_version_minor: int | float
    driver_version: str
    nvml_version: str
    gpus: list[GpuInfo]


class JobStateResources(Struct, rename="camel"):
    """Job State Resources."""

    cluster_id: PositiveInt
    gpu_ids: list[int]
    cpu_ids: list[int]
    cpu_memory_gb: int
    """does NOT include tmpfs memory (it is listed separately below)"""
    gpu_memory_gb: int
    """if less than a GB, this is rounded DOWN (so it can be 0!)"""

    slave_name: str
    slave_host: str
    slave_instance_id: str
    slave_instance_pid: int
    worker_id: int

    site_id: str
    infra_id: str
    machine_id: str

    ssh_host: str | None = None
    ssh_port: int | None = None
    ssh_username: str | None = None
    ssh_proxy_host: str | None = None
    ssh_proxy_port: int | None = None
    ssh_proxy_username: str | None = None

    port_mappings: list[JobPortMapping] = []
    gpu_details: GpuOverview | None = None

    tmpfs_mem_gb: int = 0

    @property
    def cpu_memory_mb(self) -> int:
        """CPU Memory in MB."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return self.cpu_memory_gb * 1024

    @property
    def cpu_memory_byte(self) -> int:
        """CPU Memory in bytes."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return self.cpu_memory_gb * 1024 * 1024 * 1024

    @property
    def tmpfs_mem_mb(self) -> int:
        """Tmpfs Memory in MB."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return (self.tmpfs_mem_gb or 0) * 1024

    @property
    def tmpfs_mem_byte(self) -> int:
        """Tmpfs Memory in bytes."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return (self.tmpfs_mem_gb or 0) * 1024 * 1024 * 1024

    @property
    def all_cpu_memory_gb(self) -> int:
        """All Memory in GB."""
        return self.cpu_memory_gb + (self.tmpfs_mem_gb or 0)

    @property
    def all_cpu_memory_mb(self) -> int:
        """All Memory in MB."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return (self.all_cpu_memory_gb or 0) * 1024

    @property
    def all_cpu_memory_byte(self) -> int:
        """All Memory in bytes."""
        # while GB is sometimes 1000^3 instead of 1024^3 (typically HD size etc.), for memory, base 1024 is always used.
        # (and docker seems to use 1000 based, which is silly)
        return (self.all_cpu_memory_gb or 0) * 1024 * 1024 * 1024


class JobState(Struct, rename="camel"):
    """Current Job State.

    State vs Status
       - In the english language: very related, and often synonyms
       - Typically in technical language, state is used in a more broad sense, and status is more "one-dimensional"
           -> thus you can have a state containing multiple statuses

      Here:
       - In JobState we describe the entire variable state of the Job.
         "status" holds the ID of the current discrete step in the Job's lifecycle ("the workflow of job execution")
         so "lifecycle_step" or "workflow_position" would be a synonym for our "status", but both feels too convoluted
    """

    status: JobStatus
    scheduling: JobStateScheduling
    event_times: JobEventTimes
    resources: JobStateResources | None = None
    final_usage_statistics: GPULabUsageStatistics | None = None
    waste_review: WasteReview | None = None


class SSHInfo(Struct):
    """Info about SSH Login to a resource."""

    host: str
    port: int
    username: str

    host_keys: list[str] | None
    """The public key(s) of the target SSH server, in openssh format.

    This list can be empty. It is recommended that it is not.
    The typical line you need to add to ~/.ssh/known_hosts is the host, a space, and then one of these keys.
    """

    description: str | None = None
    jump_proxy: Optional["SSHInfo"] = None

    def connect(self):
        """Create a Fabric connection."""

        return connect_with_fabric(self)


class Job(Struct, rename="camel"):
    """AI Infrastructure Service Job."""

    uuid: UUID
    slices_id: str
    name: str
    request: JobRequest
    owner: JobOwner
    state: JobState
    description: str | None = None
    restart_info: RestartInfo | None = None

    @property
    def ssh_info(self) -> list[SSHInfo]:
        """Return a list of SSHInfo objects for this job.

        These objects are compatible with the slices BI SSHInfo objects.
        """

        if (
            self.state.resources is None
            or self.state.resources.ssh_host is None
            or self.state.resources.ssh_username is None
        ):
            return []
        proxy: SSHInfo | None = None
        if self.state.resources.ssh_proxy_host is not None and self.state.resources.ssh_proxy_username is not None:
            proxy = SSHInfo(
                host=self.state.resources.ssh_proxy_host,
                port=self.state.resources.ssh_proxy_port or 22,
                username=self.state.resources.ssh_proxy_username,
                host_keys=None,
                description=None,
                jump_proxy=None,
            )
        return [
            SSHInfo(
                host=self.state.resources.ssh_host,
                port=self.state.resources.ssh_port or 22,
                username=self.state.resources.ssh_username,
                host_keys=None,
                description=None,
                jump_proxy=proxy,
            )
        ]

    @property
    def id(self) -> UUID:
        """Job ID alias."""
        return self.uuid

    @property
    def short_id(self) -> str | None:
        """Short ID that is derived from full uuid."""
        if not self.id:
            return None
        id_s = str(self.id)
        if "-" in id_s:
            return id_s[: id_s.index("-")]
        return id_s
