from __future__ import annotations

from typing import Any

from sqlalchemy import select

from app.models.base import async_session
from app.models.module_generation_project import ModuleGenerationProject
from app.services.module_async_recovery_service import (
    OBJECT_MODULE_STEP,
    TASK_MODULE_V2_VIDEO_PROMPT,
    acquire_object_lock,
    cleanup_active_if_terminal,
    ensure_object_lock_owned,
    mark_active_started,
    register_module_step_task,
    release_object_lock,
)
from app.services.module_generation_v2.config import VIDEO_PROMPT_OPTIMIZE
from app.services.module_generation_v2.flow_service import run_video_prompt_optimize_v2
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app


async def _run_video_prompt(project_id: str, step_id: str) -> None:
    lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
    if not lock_token:
        return None
    try:
        async with async_session() as db:
            result = await db.execute(
                select(ModuleGenerationProject.module)
                .where(
                    ModuleGenerationProject.id == project_id,
                    ModuleGenerationProject.deleted_at.is_(None),
                )
                .limit(1)
            )
            module = result.scalar_one_or_none()
            if not module:
                return None
            await register_module_step_task(
                module=str(module),
                project_id=project_id,
                step_id=step_id,
                step_code=VIDEO_PROMPT_OPTIMIZE,
                task_name=TASK_MODULE_V2_VIDEO_PROMPT,
            )
            await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
            await run_video_prompt_optimize_v2(
                db,
                project_id=project_id,
                step_id=step_id,
                execution_guard=lambda: ensure_object_lock_owned(token=lock_token),
            )
            await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
    finally:
        await release_object_lock(
            object_type=OBJECT_MODULE_STEP,
            object_id=step_id,
            token=lock_token,
        )


if celery_app:
    @celery_app.task(
        name="module_generation_v2.start_video_prompt_optimize",
        bind=True,
        max_retries=3,
        default_retry_delay=30,
        ignore_result=True,
    )
    def start_video_prompt_optimize_v2(self, project_id: str, step_id: str):
        try:
            run_async(_run_video_prompt(project_id, step_id))
            return None
        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_video_prompt_optimize_v2 = _DisabledTask()