from enum import Enum


class GenerationOwnerType(str, Enum):
    CHAT_GENERATION_TASK = "chat_generation_task"
    GENERATION_RECORD = "generation_record"


class GenerationMode(str, Enum):
    CHATAPI_ASYNC = "chatapi_async"
    CHATAPI_MAIN = "chatapi_main"
    CHATAPI_CHILD = "chatapi_child"
    HOT_OPENING_REPLICATE = "hot_opening_replicate"
    SHOT_REPLICATE = "shot_replicate"
    GENERATION_RECORD = "generation_record"


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


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


class ChatGenerationDisplayStatus(str, Enum):
    PENDING = "pending"
    GENERATING = "generating"
    COMPLETED = "completed"
    FAILED = "failed"
    DOWNLOAD_FAILED = "download_failed"
    DELETED = "deleted"


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"
    UPSCALE_QUEUED = "upscale_queued"
    UPSCALE_PROCESSING = "upscale_processing"
    UPSCALE_POLLING = "upscale_polling"
    UPSCALE_DOWNLOADING = "upscale_downloading"
    UPSCALE_FINALIZING = "upscale_finalizing"
    UPSCALE_RETRY_WAITING = "upscale_retry_waiting"
    UPSCALE_FAILED = "upscale_failed"
    DONE = "done"
    FAILED = "failed"
    TIMEOUT = "timeout"
    DOWNLOAD_FAILED = "download_failed"


class ChatGenerationTaskEventType(str, Enum):
    TASK_CREATED = "TASK_CREATED"
    IDEMPOTENCY_HIT = "IDEMPOTENCY_HIT"
    BATCH_CREATE_START = "BATCH_CREATE_START"
    BATCH_MAIN_CREATED = "BATCH_MAIN_CREATED"
    BATCH_CHILDREN_CREATED = "BATCH_CHILDREN_CREATED"
    BATCH_BILLING_SUCCESS = "BATCH_BILLING_SUCCESS"
    BATCH_COMMIT_SUCCESS = "BATCH_COMMIT_SUCCESS"
    CHILD_ENQUEUE_START = "CHILD_ENQUEUE_START"
    CHILD_ENQUEUE_SUCCESS = "CHILD_ENQUEUE_SUCCESS"
    CHILD_ENQUEUE_FAILED = "CHILD_ENQUEUE_FAILED"
    IMAGE_MAIN_CLAIM_ACQUIRED = "IMAGE_MAIN_CLAIM_ACQUIRED"
    IMAGE_MAIN_CLAIM_REJECTED = "IMAGE_MAIN_CLAIM_REJECTED"
    IMAGE_MAIN_CLAIM_EXPIRED = "IMAGE_MAIN_CLAIM_EXPIRED"
    IMAGE_BATCH_PROVIDER_START = "IMAGE_BATCH_PROVIDER_START"
    IMAGE_BATCH_PROVIDER_SUCCESS = "IMAGE_BATCH_PROVIDER_SUCCESS"
    IMAGE_BATCH_PROVIDER_FAILED = "IMAGE_BATCH_PROVIDER_FAILED"
    IMAGE_BATCH_SPLIT_START = "IMAGE_BATCH_SPLIT_START"
    IMAGE_BATCH_SPLIT_SUCCESS = "IMAGE_BATCH_SPLIT_SUCCESS"
    IMAGE_BATCH_SPLIT_FAILED = "IMAGE_BATCH_SPLIT_FAILED"
    MAIN_STATUS_AGGREGATED = "MAIN_STATUS_AGGREGATED"
    CHILD_RESOURCE_DELETE_START = "CHILD_RESOURCE_DELETE_START"
    CHILD_RESOURCE_DELETE_SUCCESS = "CHILD_RESOURCE_DELETE_SUCCESS"
    BATCH_GROUP_DELETE_SUCCESS = "BATCH_GROUP_DELETE_SUCCESS"
    BATCH_RECOVERY_RECONCILED = "BATCH_RECOVERY_RECONCILED"

    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_SCHEDULED = "POLL_SCHEDULED"
    POLL_SKIP_NOT_DUE = "POLL_SKIP_NOT_DUE"
    POLL_DISPATCH_DUE = "POLL_DISPATCH_DUE"
    POLL_DISPATCH_SKIP = "POLL_DISPATCH_SKIP"
    POLL_SUCCESS = "POLL_SUCCESS"
    POLL_FAILED = "POLL_FAILED"
    POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY = "POLL_SUCCESS_AFTER_TIMEOUT_RECOVERY"
    FINAL_POLL_BEFORE_TIMEOUT = "FINAL_POLL_BEFORE_TIMEOUT"
    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"
    GENERATION_RECOVERY_TIMEOUT = "GENERATION_RECOVERY_TIMEOUT"

    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"

    UPSCALE_SNAPSHOT_MATCHED = "UPSCALE_SNAPSHOT_MATCHED"
    UPSCALE_SNAPSHOT_BYPASSED = "UPSCALE_SNAPSHOT_BYPASSED"
    UPSCALE_SOURCE_READY = "UPSCALE_SOURCE_READY"
    UPSCALE_ENQUEUE = "UPSCALE_ENQUEUE"
    UPSCALE_START = "UPSCALE_START"
    UPSCALE_REMOTE_SUBMIT = "UPSCALE_REMOTE_SUBMIT"
    UPSCALE_REMOTE_POLL = "UPSCALE_REMOTE_POLL"
    UPSCALE_REMOTE_RESULT_READY = "UPSCALE_REMOTE_RESULT_READY"
    UPSCALE_RETRY_WAITING = "UPSCALE_RETRY_WAITING"
    UPSCALE_SUCCESS = "UPSCALE_SUCCESS"
    UPSCALE_FAILED = "UPSCALE_FAILED"
    UPSCALE_RECOVERY_ENQUEUE = "UPSCALE_RECOVERY_ENQUEUE"

    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"
    STALE_ATTEMPT_MESSAGE_SKIPPED = "STALE_ATTEMPT_MESSAGE_SKIPPED"
    GENERATION_RECORD_ENQUEUE_SUCCESS = "GENERATION_RECORD_ENQUEUE_SUCCESS"
    GENERATION_RECORD_ENQUEUE_FAILED = "GENERATION_RECORD_ENQUEUE_FAILED"

    TASK_TIMEOUT = "TASK_TIMEOUT"
    TASK_FAILED = "TASK_FAILED"


CHAT_TOP_LEVEL_MODES = {
    GenerationMode.CHATAPI_ASYNC.value,
    GenerationMode.CHATAPI_MAIN.value,
}

CHAT_RESOURCE_MODES = {
    GenerationMode.CHATAPI_ASYNC.value,
    GenerationMode.CHATAPI_CHILD.value,
}

CHAT_EXECUTABLE_MODES = {
    GenerationMode.CHATAPI_ASYNC.value,
    GenerationMode.CHATAPI_CHILD.value,
}

ALLOWED_GENERATION_MODES = {
    *CHAT_EXECUTABLE_MODES,
    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,
    ChatGenerationPipelineStage.UPSCALE_FAILED.value,
}

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

PROVIDER_SUCCESS_STATUSES = {"succeeded", "success", "completed", "done"}
PROVIDER_FAILED_STATUSES = {"failed", "error", "canceled", "cancelled"}