from __future__ import annotations

import uuid
from collections.abc import Awaitable, Callable
from typing import Any

from app.config import settings
from app.enums.celery_queue import CeleryQueue, CeleryTaskName
from app.enums.celery_runtime import CeleryRuntimeDomain
from app.models.base import async_session
from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
from app.services.redis_registry_service import RedisExecutionLockLease
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_with_execution_lock(
    upscale_task_id: str,
    callback: Callable[[str, str, Callable[[], Awaitable[None]]], Awaitable[None]],
    *,
    task_name: str,
    queue: str,
    pipeline_stage: str,
) -> None:
    lock_key = f"{settings.VIDEO_UPSCALE_EXECUTION_LOCK_KEY_PREFIX}:{upscale_task_id}"
    token = uuid.uuid4().hex
    lease = await CeleryRuntimeLease.acquire(
        identity=RuntimeIdentity(
            domain=CeleryRuntimeDomain.VIDEO_UPSCALE.value,
            owner_type="video_upscale_task",
            owner_id=upscale_task_id,
            attempt_no=1,
            task_name=task_name,
            queue=queue,
            registry_item_id=f"video_upscale:{upscale_task_id}",
        ),
        lock_key=lock_key,
        hash_key=settings.VIDEO_UPSCALE_ACTIVE_REDIS_HASH_KEY,
        zset_key=settings.VIDEO_UPSCALE_ACTIVE_REDIS_ZSET_KEY,
        token=token,
        ttl_seconds=max(30, int(settings.VIDEO_UPSCALE_EXECUTION_LOCK_TTL_SECONDS or 900)),
        heartbeat_interval_seconds=max(1, int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 20)),
        pipeline_stage=pipeline_stage,
    )
    if lease is None:
        # 重复消息已有其他 Worker 推进，不属于业务失败。
        return
    try:
        await callback(upscale_task_id, lease.token, lease.ensure_owned)
        await lease.ensure_owned()
    finally:
        await lease.close()


async def _run_local(upscale_task_id: str) -> None:
    async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
        async with async_session() as db:
            await run_local_upscale(db, task_id, execution_token=token, execution_guard=guard)
    await _run_with_execution_lock(
        upscale_task_id,
        _execute,
        task_name=CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value,
        queue=CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value,
        pipeline_stage="execute_local",
    )


async def _run_submit(upscale_task_id: str, *, count_attempt: bool = True) -> None:
    async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
        async with async_session() as db:
            await run_remote_submit(
                db, task_id, count_attempt=count_attempt, execution_token=token, execution_guard=guard
            )
    await _run_with_execution_lock(
        upscale_task_id,
        _execute,
        task_name=CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value,
        queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
        pipeline_stage="submit_remote",
    )


async def _run_poll(upscale_task_id: str) -> None:
    async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
        async with async_session() as db:
            await run_remote_poll(db, task_id, execution_token=token, execution_guard=guard)
    await _run_with_execution_lock(
        upscale_task_id,
        _execute,
        task_name=CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value,
        queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
        pipeline_stage="poll_remote",
    )


async def _run_download(upscale_task_id: str) -> None:
    async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
        async with async_session() as db:
            await run_remote_result_download(db, task_id, execution_token=token, execution_guard=guard)
    await _run_with_execution_lock(
        upscale_task_id,
        _execute,
        task_name=CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value,
        queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
        pipeline_stage="download_remote_result",
    )


async def _run_finalize(upscale_task_id: str) -> None:
    async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
        async with async_session() as db:
            await run_finalize_upscale(db, task_id, execution_token=token, execution_guard=guard)
    await _run_with_execution_lock(
        upscale_task_id,
        _execute,
        task_name=CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value,
        queue=CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value,
        pipeline_stage="finalize",
    )


async def _run_recovery() -> dict[str, Any]:
    lease = await RedisExecutionLockLease.acquire(
        lock_key=settings.VIDEO_UPSCALE_RECOVERY_LOCK_KEY,
        ttl_seconds=max(30, int(settings.VIDEO_UPSCALE_EXECUTION_LOCK_TTL_SECONDS or 900)),
        log_context="video_upscale_recovery",
        renew_interval_seconds=max(1, int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 20)),
    )
    if lease is None:
        return {"checked": 0, "results": {"lock_busy": 1}}
    try:
        async with async_session() as db:
            result = await recover_video_upscale_tasks_once(db)
        await lease.ensure_owned()
        return result
    finally:
        await lease.close()


if celery_app:

    @celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value, 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=CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value, 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=CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value, 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=CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value, 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=CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value, 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=CeleryTaskName.VIDEO_UPSCALE_RECOVER.value, bind=True, max_retries=2)
    def recover_once(self) -> dict[str, Any]:
        try:
            return run_async(_run_recovery())
        except Exception as exc:
            raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))

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()