from __future__ import annotations

from typing import Any

from app.config import settings
from app.models.base import async_session
from app.services.video_upscale.task_service import (
    recover_video_upscale_tasks_once,
    run_finalize_upscale,
    run_local_upscale,
    run_remote_poll,
    run_remote_result_download,
    run_remote_submit,
)
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app


async def _run_local(upscale_task_id: str) -> None:
    async with async_session() as db:
        await run_local_upscale(db, upscale_task_id)


async def _run_submit(upscale_task_id: str, *, count_attempt: bool = True) -> None:
    async with async_session() as db:
        await run_remote_submit(db, upscale_task_id, count_attempt=count_attempt)


async def _run_poll(upscale_task_id: str) -> None:
    async with async_session() as db:
        await run_remote_poll(db, upscale_task_id)


async def _run_download(upscale_task_id: str) -> None:
    async with async_session() as db:
        await run_remote_result_download(db, upscale_task_id)


async def _run_finalize(upscale_task_id: str) -> None:
    async with async_session() as db:
        await run_finalize_upscale(db, upscale_task_id)


async def _run_recovery() -> dict[str, Any]:
    async with async_session() as db:
        return await recover_video_upscale_tasks_once(db)


if celery_app:

    @celery_app.task(name="video_upscale.execute_local", bind=True, max_retries=2)
    def execute_local(self, upscale_task_id: str) -> None:
        try:
            return run_async(_run_local(upscale_task_id))
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))


    @celery_app.task(name="video_upscale.submit_remote", bind=True, max_retries=2)
    def submit_remote(self, upscale_task_id: str, count_attempt: bool = True) -> None:
        try:
            return run_async(_run_submit(upscale_task_id, count_attempt=count_attempt))
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))


    @celery_app.task(name="video_upscale.poll_remote", bind=True, max_retries=2)
    def poll_remote(self, upscale_task_id: str) -> None:
        try:
            return run_async(_run_poll(upscale_task_id))
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))


    @celery_app.task(name="video_upscale.download_remote_result", bind=True, max_retries=2)
    def download_remote_result(self, upscale_task_id: str) -> None:
        try:
            return run_async(_run_download(upscale_task_id))
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))


    @celery_app.task(name="video_upscale.finalize", bind=True, max_retries=2)
    def finalize(self, upscale_task_id: str) -> None:
        try:
            return run_async(_run_finalize(upscale_task_id))
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))


    @celery_app.task(name="video_upscale.recover_once", bind=True)
    def recover_once(self) -> dict[str, Any]:
        return run_async(_run_recovery())

else:

    class _DisabledTask:
        def delay(self, *args: Any, **kwargs: Any) -> None:
            raise RuntimeError("Celery is disabled")

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

    execute_local = _DisabledTask()
    submit_remote = _DisabledTask()
    poll_remote = _DisabledTask()
    download_remote_result = _DisabledTask()
    finalize = _DisabledTask()
    recover_once = _DisabledTask()