from enum import Enum


class GenerationMode(str, Enum):
    CHATAPI_ASYNC = "chatapi_async"
    HOT_OPENING_REPLICATE = "hot_opening_replicate"
    SHOT_REPLICATE = "shot_replicate"


class GenerationType(str, Enum):
    IMAGE = "image"
    VIDEO = "video"


class ChatGenerationTaskStatus(str, Enum):
    PENDING = "pending"
    GENERATING = "generating"
    COMPLETED = "completed"
    FAILED = "failed"


class ChatGenerationPipelineStage(str, Enum):
    QUEUED = "queued"
    PREPARING = "preparing"
    CREATING_PROVIDER_TASK = "creating_provider_task"
    WAITING_REMOTE = "waiting_remote"
    POLLING = "polling"
    RESULT_READY = "result_ready"
    DOWNLOAD_QUEUED = "download_queued"
    DOWNLOADING = "downloading"
    RETRY_WAITING = "retry_waiting"
    DONE = "done"
    FAILED = "failed"
    TIMEOUT = "timeout"
    DOWNLOAD_FAILED = "download_failed"


class ChatGenerationTaskEventType(str, Enum):
    PROMPT_CONCAT_START = "PROMPT_CONCAT_START"
    PROMPT_CONCAT_SUCCESS = "PROMPT_CONCAT_SUCCESS"

    PROVIDER_CREATE_START = "PROVIDER_CREATE_START"
    PROVIDER_CREATE_SUCCESS = "PROVIDER_CREATE_SUCCESS"
    PROVIDER_CREATE_FAILED = "PROVIDER_CREATE_FAILED"

    POLL_START = "POLL_START"
    POLL_PENDING = "POLL_PENDING"
    POLL_SUCCESS = "POLL_SUCCESS"
    POLL_FAILED = "POLL_FAILED"
    POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY = "POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY"
    FINAL_POLL_BEFORE_TIMEOUT_ERROR = "FINAL_POLL_BEFORE_TIMEOUT_ERROR"
    FINAL_POLL_BEFORE_TIMEOUT_PENDING = "FINAL_POLL_BEFORE_TIMEOUT_PENDING"
    GENERATION_RECOVERY_ENQUEUE = "GENERATION_RECOVERY_ENQUEUE"

    DOWNLOAD_ENQUEUE = "DOWNLOAD_ENQUEUE"
    DOWNLOAD_ENQUEUE_FAILED = "DOWNLOAD_ENQUEUE_FAILED"
    DOWNLOAD_START = "DOWNLOAD_START"
    DOWNLOAD_SUCCESS = "DOWNLOAD_SUCCESS"
    DOWNLOAD_RETRY_WAITING = "DOWNLOAD_RETRY_WAITING"
    DOWNLOAD_RETRY_ENQUEUE = "DOWNLOAD_RETRY_ENQUEUE"
    DOWNLOAD_RETRY_ENQUEUE_FAILED = "DOWNLOAD_RETRY_ENQUEUE_FAILED"
    DOWNLOAD_RETRY_NOT_DUE = "DOWNLOAD_RETRY_NOT_DUE"
    DOWNLOAD_STUCK_RECOVER = "DOWNLOAD_STUCK_RECOVER"
    DOWNLOAD_RECOVERY_ENQUEUE = "DOWNLOAD_RECOVERY_ENQUEUE"
    DOWNLOAD_RECOVERY_ENQUEUE_FAILED = "DOWNLOAD_RECOVERY_ENQUEUE_FAILED"
    DOWNLOAD_FAILED = "DOWNLOAD_FAILED"
    DOWNLOAD_FAILED_NON_RETRYABLE = "DOWNLOAD_FAILED_NON_RETRYABLE"

    DOWNLOAD_SKIP_TASK_MISSING = "DOWNLOAD_SKIP_TASK_MISSING"
    DOWNLOAD_SKIP_INVALID_MODE = "DOWNLOAD_SKIP_INVALID_MODE"
    DOWNLOAD_SKIP_NOT_GENERATING = "DOWNLOAD_SKIP_NOT_GENERATING"
    DOWNLOAD_SKIP_ALREADY_COMPLETED = "DOWNLOAD_SKIP_ALREADY_COMPLETED"
    DOWNLOAD_SKIP_NO_REMOTE_RESULT_URL = "DOWNLOAD_SKIP_NO_REMOTE_RESULT_URL"
    DOWNLOAD_SKIP_STAGE_NOT_ALLOWED = "DOWNLOAD_SKIP_STAGE_NOT_ALLOWED"
    DOWNLOAD_SKIP_DOWNLOADING_LEASE_ALIVE = "DOWNLOAD_SKIP_DOWNLOADING_LEASE_ALIVE"
    DOWNLOAD_SKIP_RETRY_NOT_DUE = "DOWNLOAD_SKIP_RETRY_NOT_DUE"
    DOWNLOAD_SKIP_FINAL_STATE = "DOWNLOAD_SKIP_FINAL_STATE"
    DOWNLOAD_SKIP_DISABLED = "DOWNLOAD_SKIP_DISABLED"

    TASK_TIMEOUT = "TASK_TIMEOUT"


ALLOWED_GENERATION_MODES = {
    GenerationMode.CHATAPI_ASYNC.value,
    GenerationMode.HOT_OPENING_REPLICATE.value,
    GenerationMode.SHOT_REPLICATE.value,
}

FINAL_CHAT_GENERATION_STAGES = {
    ChatGenerationPipelineStage.DONE.value,
    ChatGenerationPipelineStage.FAILED.value,
    ChatGenerationPipelineStage.TIMEOUT.value,
    ChatGenerationPipelineStage.DOWNLOAD_FAILED.value,
}

DOWNLOAD_RECOVERABLE_STAGES = {
    ChatGenerationPipelineStage.RESULT_READY.value,
    ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value,
    ChatGenerationPipelineStage.DOWNLOADING.value,
    ChatGenerationPipelineStage.RETRY_WAITING.value,
}