from __future__ import annotations
from enum import IntEnum

from ..protogen.common.task_status_pb2 import (
    TASK_STATUS_UNSPECIFIED,
    TASK_STATUS_CREATING,
    TASK_STATUS_SUBMITTED,
    TASK_STATUS_DISPATCHED,
    TASK_STATUS_COMPLETED,
    TASK_STATUS_ERROR,
    TASK_STATUS_TIMEOUT,
    TASK_STATUS_CANCELLING,
    TASK_STATUS_CANCELLED,
    TASK_STATUS_PROCESSING,
    TASK_STATUS_PROCESSED,
    TASK_STATUS_RETRIED,
    TASK_STATUS_PENDING,
    TASK_STATUS_PAUSED,
)
from ..protogen.common.events_common_pb2 import (
    EVENTS_ENUM_UNSPECIFIED,
    EVENTS_ENUM_NEW_TASK,
    EVENTS_ENUM_TASK_STATUS_UPDATE,
    EVENTS_ENUM_NEW_RESULT,
    EVENTS_ENUM_RESULT_STATUS_UPDATE,
    EVENTS_ENUM_RESULT_OWNER_UPDATE,
)
from ..protogen.common.session_status_pb2 import (
    SessionStatus as RawSessionStatus,
    _SESSIONSTATUS,
    SESSION_STATUS_UNSPECIFIED,
    SESSION_STATUS_RUNNING,
    SESSION_STATUS_CANCELLED,
    SESSION_STATUS_PAUSED,
    SESSION_STATUS_CLOSED,
    SESSION_STATUS_PURGED,
    SESSION_STATUS_DELETED,
)
from ..protogen.common.result_status_pb2 import (
    ResultStatus as RawResultStatus,
    _RESULTSTATUS,
    RESULT_STATUS_UNSPECIFIED,
    RESULT_STATUS_CREATED,
    RESULT_STATUS_COMPLETED,
    RESULT_STATUS_ABORTED,
    RESULT_STATUS_DELETED,
    RESULT_STATUS_NOTFOUND,
)
from ..protogen.common.health_checks_common_pb2 import (
    HEALTH_STATUS_ENUM_UNSPECIFIED,
    HEALTH_STATUS_ENUM_HEALTHY,
    HEALTH_STATUS_ENUM_DEGRADED,
    HEALTH_STATUS_ENUM_UNHEALTHY,
)
from ..protogen.common.worker_common_pb2 import HealthCheckReply
from ..protogen.common.sort_direction_pb2 import SORT_DIRECTION_ASC, SORT_DIRECTION_DESC

# This file is necessary because the grpc types aren't considered proper types


class TaskStatus(IntEnum):
    UNSPECIFIED = TASK_STATUS_UNSPECIFIED
    CREATING = TASK_STATUS_CREATING
    SUBMITTED = TASK_STATUS_SUBMITTED
    DISPATCHED = TASK_STATUS_DISPATCHED
    COMPLETED = TASK_STATUS_COMPLETED
    ERROR = TASK_STATUS_ERROR
    TIMEOUT = TASK_STATUS_TIMEOUT
    CANCELLING = TASK_STATUS_CANCELLING
    CANCELLED = TASK_STATUS_CANCELLED
    PROCESSING = TASK_STATUS_PROCESSING
    PROCESSED = TASK_STATUS_PROCESSED
    RETRIED = TASK_STATUS_RETRIED
    PENDING = TASK_STATUS_PENDING
    PAUSED = TASK_STATUS_PAUSED


class EventTypes(IntEnum):
    UNSPECIFIED = EVENTS_ENUM_UNSPECIFIED
    NEW_TASK = EVENTS_ENUM_NEW_TASK
    TASK_STATUS_UPDATE = EVENTS_ENUM_TASK_STATUS_UPDATE
    NEW_RESULT = EVENTS_ENUM_NEW_RESULT
    RESULT_STATUS_UPDATE = EVENTS_ENUM_RESULT_STATUS_UPDATE
    RESULT_OWNER_UPDATE = EVENTS_ENUM_RESULT_OWNER_UPDATE

    @classmethod
    def from_string(cls, name: str):
        return getattr(cls, name.upper())


class SessionStatus(IntEnum):
    @staticmethod
    def name_from_value(status: RawSessionStatus) -> str:
        return _SESSIONSTATUS.values_by_number[status].name

    UNSPECIFIED = SESSION_STATUS_UNSPECIFIED
    RUNNING = SESSION_STATUS_RUNNING
    CANCELLED = SESSION_STATUS_CANCELLED
    PAUSED = SESSION_STATUS_PAUSED
    CLOSED = SESSION_STATUS_CLOSED
    PURGED = SESSION_STATUS_PURGED
    DELETED = SESSION_STATUS_DELETED


class ResultStatus(IntEnum):
    @staticmethod
    def name_from_value(status: RawResultStatus) -> str:
        return _RESULTSTATUS.values_by_number[status].name

    UNSPECIFIED = RESULT_STATUS_UNSPECIFIED
    CREATED = RESULT_STATUS_CREATED
    COMPLETED = RESULT_STATUS_COMPLETED
    ABORTED = RESULT_STATUS_ABORTED
    DELETED = RESULT_STATUS_DELETED
    NOTFOUND = RESULT_STATUS_NOTFOUND


class ServiceHealthCheckStatus(IntEnum):
    UNSPECIFIED = HEALTH_STATUS_ENUM_UNSPECIFIED
    HEALTHY = HEALTH_STATUS_ENUM_HEALTHY
    DEGRADED = HEALTH_STATUS_ENUM_DEGRADED
    UNHEALTHY = HEALTH_STATUS_ENUM_UNHEALTHY


class HealthCheckStatus(IntEnum):
    UNKNOWN = HealthCheckReply.UNKNOWN
    SERVING = HealthCheckReply.SERVING
    NOT_SERVING = HealthCheckReply.NOT_SERVING


class Direction(IntEnum):
    ASC = SORT_DIRECTION_ASC
    DESC = SORT_DIRECTION_DESC
