from __future__ import annotations

from typing import Any

from app.enums.hot_opening_replicate import HotOpeningStepCodeEnum, ModuleCodeEnum
from app.models.base import async_session
from app.services.hot_opening_replicate_service import run_image_prompt_optimize, run_video_prompt_optimize
from app.services.module_async_recovery_service import (
    OBJECT_MODULE_STEP,
    TASK_HOT_IMAGE_PROMPT,
    TASK_HOT_VIDEO_PROMPT,
    acquire_object_lock,
    cleanup_active_if_terminal,
    mark_active_started,
    release_object_lock,
    register_module_step_task,
)
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app

MODULE = ModuleCodeEnum.HOT_OPENING_REPLICATE.value


async def _run_image_prompt(project_id: str, step_id: str | None = None):
    lock_token: str | None = None
    if step_id:
        lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
        if not lock_token:
            # 同一个 step 的重复消息正在被其它 worker 处理：当前消息直接忽略，active registry 保留等待真正任务完成或 lease 过期。
            return None
        await register_module_step_task(
            module=MODULE,
            project_id=project_id,
            step_id=step_id,
            step_code=HotOpeningStepCodeEnum.IMAGE_PROMPT_OPTIMIZE.value,
            task_name=TASK_HOT_IMAGE_PROMPT,
        )
        await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
    try:
        async with async_session() as db:
            result = await run_image_prompt_optimize(db, project_id=project_id, step_id=step_id)
            await db.commit()
            if step_id:
                await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
            return result
    finally:
        if step_id:
            await release_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id, token=lock_token)


async def _run_video_prompt(project_id: str, step_id: str | None = None):
    lock_token: str | None = None
    if step_id:
        lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
        if not lock_token:
            return None
        await register_module_step_task(
            module=MODULE,
            project_id=project_id,
            step_id=step_id,
            step_code=HotOpeningStepCodeEnum.VIDEO_PROMPT_OPTIMIZE.value,
            task_name=TASK_HOT_VIDEO_PROMPT,
        )
        await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
    try:
        async with async_session() as db:
            result = await run_video_prompt_optimize(db, project_id=project_id, step_id=step_id)
            await db.commit()
            if step_id:
                await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
            return result
    finally:
        if step_id:
            await release_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id, token=lock_token)


if celery_app:
    @celery_app.task(name="hot_opening.start_image_prompt_optimize", bind=True, max_retries=3, default_retry_delay=30)
    def start_image_prompt_optimize(self, project_id: str, step_id: str | None = None):
        """手动触发后的图片 AI 提词任务。

        max_retries 只有在显式 self.retry() 时才会生效。
        这里仅兜底 Celery 包装层 / DB commit / Redis registry 等未被 service 捕获的异常；
        service 内部已经落库为业务失败的情况不会抛出异常，也不会重复 retry。
        """
        try:
            return run_async(_run_image_prompt(project_id, step_id))
        except Exception as exc:
            raise self.retry(exc=exc) from exc

    @celery_app.task(name="hot_opening.start_video_prompt_optimize", bind=True, max_retries=3, default_retry_delay=30)
    def start_video_prompt_optimize(self, project_id: str, step_id: str | None = None):
        """手动触发后的视频 AI 提词任务。

        max_retries 只有在显式 self.retry() 时才会生效。
        这里仅兜底 Celery 包装层 / DB commit / Redis registry 等未被 service 捕获的异常；
        service 内部已经落库为业务失败的情况不会抛出异常，也不会重复 retry。
        """
        try:
            return run_async(_run_video_prompt(project_id, step_id))
        except Exception as exc:
            raise self.retry(exc=exc) from exc
else:
    class _DisabledTask:
        def delay(self, *args: Any, **kwargs: Any):
            raise RuntimeError("Celery is disabled")

        def apply_async(self, *args: Any, **kwargs: Any):
            raise RuntimeError("Celery is disabled")

    start_image_prompt_optimize = _DisabledTask()
    start_video_prompt_optimize = _DisabledTask()