from enum import Enum
from typing import Union
from .jobs.nbox_ws_pb2 import JobRequest, ListJobsRequest, UpdateJobRequest
from .deploy.serve_pb2 import ServingListRequest, ServingRequest, UpdateServingRequest, ModelListRequest, ModelRequest

# To be used as TYPE.ACTION_[RESOURCE]_STATUS
class Actions(Enum):
    # Common actions
    INIT = "INIT_DONE"
    # CODE_UPLOADED = "UPLOAD_CODE_DONE"
    UPDATED_STATUS = "UPDATE_STATUS_DONE"
    UPDATE_RESOURCES_INIT = "UPDATE_RESOURCES_INIT"
    UPDATE_RESOURCES_DONE = "UPDATE_RESOURCES_DONE"
    UPDATE_RESOURCES_FAIL = "UPDATE_RESOURCES_FAIL"
    UPDATE_INFO_INIT = "UPDATE_INFO_INIT"
    UPDATE_INFO_DONE = "UPDATE_INFO_DONE"
    UPDATE_INFO_FAIL = "UPDATE_INFO_FAIL"
    SOURCE_CODE_UPDATE_DONE = "UPDATE_SOURCE_CODE_DONE"
    SOURCE_CODE_UPDATE_FAIL = "UPDATE_SOURCE_CODE_FAIL"
    DELETE_INIT = "DELETE_INIT"
    DELETE_DONE = "DELETE_DONE"
    DELETE_FAIL = "DELETE_FAIL"
    UPDATE_FEATURE_GATES_INIT = "UPDATE_FEATURE_GATES_INIT"
    UPDATE_FEATURE_GATES_DONE = "UPDATE_FEATURE_GATES_DONE"
    UPDATE_FEATURE_GATES_FAIL = "UPDATE_FEATURE_GATES_FAIL"
    UPDATE_AUTH_INFO_INIT = "UPDATE_AUTH_INFO_INIT"
    UPDATE_AUTH_INFO_DONE = "UPDATE_AUTH_INFO_DONE"
    UPDATE_AUTH_INFO_FAIL = "UPDATE_AUTH_INFO_FAIL"
    K8S_CREATE_INIT = "K8S_CREATE_INIT"
    K8S_CREATE_DONE = "K8S_CREATE_DONE"
    K8S_CREATE_FAIL = "K8S_CREATE_FAIL"

    PAUSE_INIT = "PAUSE_INIT"
    PAUSE_DONE = "PAUSE_DONE"
    PAUSE_FAIL = "PAUSE_FAIL"
    RESUME_INIT = "RESUME_INIT"
    RESUME_DONE = "RESUME_DONE"
    RESUME_FAIL = "RESUME_FAIL"
    RESCHEDULE_INIT = "RESCHEDULE_INIT"
    RESCHEDULE_DONE = "RESCHEDULE_DONE"
    RESCHEDULE_FAIL = "RESCHEDULE_FAIL"
    TRIGGER_INIT = "TRIGGER_INIT"
    TRIGGER_DONE = "TRIGGER_DONE"
    TRIGGER_FAIL = "TRIGGER_FAIL"
    PREPARE_RUN_INIT = "PREPARE_RUN_INIT"
    PREPARE_RUN_DONE = "PREPARE_RUN_DONE"
    PREPARE_RUN_FAIL = "PREPARE_RUN_FAIL"

    # Serving Actions
    UNSPECIFIED = "UNSPECIFIED"
    PIN_INIT = "PIN_INIT"
    PIN_DONE = "PIN_DONE"
    PIN_FAIL = "PIN_FAIL"
    UNPIN_INIT = "UNPIN_INIT"
    UNPIN_DONE = "UNPIN_DONE"
    UNPIN_FAIL = "UNPIN_FAIL"
    UPDATE_SCALE_INIT = "UPDATE_SCALE_INIT"
    UPDATE_SCALE_DONE = "UPDATE_SCALE_DONE"
    UPDATE_SCALE_FAIL = "UPDATE_SCALE_FAIL"
    UPDATE_ENDPOINT_ACCESS_KEY_INIT = "UPDATE_ENDPOINT_ACCESS_KEY_INIT"
    UPDATE_ENDPOINT_ACCESS_KEY_DONE = "UPDATE_ENDPOINT_ACCESS_KEY_DONE"
    UPDATE_ENDPOINT_ACCESS_KEY_FAIL = "UPDATE_ENDPOINT_ACCESS_KEY_FAIL"

    @classmethod
    def from_string(cls, value: str) -> "Actions":
        try:
            return cls(value)
        except ValueError:
            return cls.UNSPECIFIED


class Type(Enum):
    JOB = "JOB"
    SERVING = "SERVING"
    MODEL = "MODEL"


JobsVersionType = dict[str, Union[str, list[str]]]


OneClickRequestTypes = Union[
    JobRequest,
    ListJobsRequest,
    UpdateJobRequest,
    ServingListRequest,
    ServingRequest,
    UpdateServingRequest,
    ModelRequest,
    ModelListRequest,
]
