# app/tasks/generation_recovery_tasks.py
from __future__ import annotations

import logging
from typing import Any, Awaitable, Callable, Dict

from app.config import settings
from app.enums.celery_queue import CeleryQueue
from app.models.base import async_session
from app.services.redis_registry_service import get_registry_redis, redis_acquire_lock, redis_release_lock
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app

logger = logging.getLogger("video_gen")

RECOVERY_QUEUE = settings.CELERY_RECOVERY_QUEUE or CeleryQueue.GEN_RECOVERY.value
RecoveryRunner = Callable[[], Awaitable[Dict[str, Any]]]


async def _run_download_once() -> Dict[str, Any]:
    from app.services.generation.recovery_service import recover_download_tasks_once

    async with async_session() as db:
        return await recover_download_tasks_once(db)


async def _run_generation_once() -> Dict[str, Any]:
    from app.services.generation.recovery_service import recover_generation_tasks_once

    async with async_session() as db:
        return await recover_generation_tasks_once(db)


async def _run_due_poll_dispatch_once() -> Dict[str, Any]:
    from app.services.generation.recovery_service import dispatch_due_poll_tasks_once

    async with async_session() as db:
        return await dispatch_due_poll_tasks_once(db)


async def _run_module_async_once() -> Dict[str, Any]:
    from app.services.module_async_recovery_service import recover_module_async_tasks_once

    async with async_session() as db:
        return await recover_module_async_tasks_once(db)


async def _run_shot_split_once() -> Dict[str, Any]:
    from app.services.shot_replicate_recovery_service import recover_shot_split_tasks_once

    async with async_session() as db:
        return await recover_shot_split_tasks_once(db)


async def _run_video_upscale_once() -> Dict[str, Any]:
    from app.services.video_upscale.task_service import recover_video_upscale_tasks_once

    async with async_session() as db:
        return await recover_video_upscale_tasks_once(db)


async def _run_with_execution_lock(
    *,
    lock_key: str,
    log_context: str,
    runner: RecoveryRunner,
    ttl_seconds: int | None = None,
) -> Dict[str, Any]:
    """恢复任务执行锁。

    worker_ready 的启动锁只保证“只投递一次”；如果 broker 中残留旧消息，
    或者人工手动触发恢复任务，仍可能并发执行。这里再加执行锁，避免多个
    恢复扫描同时扫库、抢行锁、抢连接池。
    """
    redis = await get_registry_redis()
    token: str | None = None
    if redis is not None:
        token = await redis_acquire_lock(
            lock_key=lock_key,
            ttl_seconds=int(ttl_seconds or settings.CELERY_RECOVERY_TASK_LOCK_TTL_SECONDS or 600),
            log_context=log_context,
        )
        if not token:
            return {"skipped": "lock_held", "lock_key": lock_key}
    else:
        # Redis 不可用时仍允许 DB fallback 执行一次，避免恢复能力彻底失效。
        logger.warning("恢复任务执行锁不可用，降级直接执行。context=%s", log_context)

    try:
        result = await runner()
        result["execution_lock"] = "lock_acquired" if token else "redis_unavailable_run_db_fallback"
        return result
    finally:
        if token:
            await redis_release_lock(lock_key=lock_key, token=token, log_context=log_context)


async def _is_lock_held(lock_key: str) -> bool:
    redis = await get_registry_redis()
    if redis is None:
        return False
    try:
        return bool(await redis.exists(lock_key))
    except Exception:
        return False


async def _startup_or_generation_recovery_running() -> str | None:
    # Beat 触发 dispatcher 时，如果启动容灾或完整生成容灾还在跑，直接跳过本轮。
    # gen_recovery concurrency=1 已经能串行；这里是多机部署、残留消息、手动触发时的双保险。
    lock_checks = [
        ("startup_recovery", settings.CELERY_RECOVERY_STARTUP_TASK_LOCK_KEY),
        ("generation_recovery", settings.GENERATION_RECOVERY_LOCK_KEY),
    ]
    for name, lock_key in lock_checks:
        if await _is_lock_held(lock_key):
            return name
    return None


async def _run_due_poll_dispatch_with_guard() -> Dict[str, Any]:
    running = await _startup_or_generation_recovery_running()
    if running:
        return {"skipped": "recovery_lock_held", "lock": running}
    return await _run_due_poll_dispatch_once()


async def _acquire_download_recovery_loop_lock() -> tuple[bool, str]:
    """下载恢复循环锁。

    Redis 不可用时降级为直接执行 DB fallback，避免恢复能力彻底失效；
    Redis 可用但锁被其他 worker 持有时，本轮跳过，不再重复投递下一轮。
    """
    redis = await get_registry_redis()
    if redis is None:
        return True, "redis_unavailable_run_db_fallback"
    token = await redis_acquire_lock(
        lock_key=settings.DOWNLOAD_RECOVERY_LOOP_LOCK_KEY,
        ttl_seconds=int(settings.DOWNLOAD_RECOVERY_LOOP_LOCK_TTL_SECONDS or 55),
        log_context="download_recovery_loop",
    )
    return (bool(token), "lock_acquired" if token else "lock_held")


def _schedule_next_download_recovery_loop() -> None:
    if not celery_app or not bool(getattr(settings, "DOWNLOAD_RECOVERY_LOOP_ENABLED", False)):
        return
    try:
        recover_download_tasks_once.apply_async(
            countdown=max(1, int(settings.DOWNLOAD_RECOVERY_INTERVAL_SECONDS or 60)),
            queue=RECOVERY_QUEUE,
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )
    except Exception:
        logger.exception("下载恢复循环下一轮投递失败")


async def _run_startup_recovery_once() -> Dict[str, Any]:
    """启动容灾协调器：串行跑恢复扫描。

    真实业务任务仍投递回原队列：
    - 创建/提词/视频分析 -> gen_chatapi_create
    - provider poll -> gen_provider_poll
    - 下载/ffmpeg 切片 -> gen_result_download
    - 本地视频超分 -> gen_video_upscale_local
    - 火山视频超分 -> gen_video_upscale_remote
    恢复扫描本身只走 gen_recovery，避免堵住业务 worker。
    """
    return await _run_with_execution_lock(
        lock_key=settings.CELERY_RECOVERY_STARTUP_TASK_LOCK_KEY,
        log_context="startup_recovery_once",
        runner=_run_startup_recovery_steps,
    )


async def _run_startup_recovery_steps() -> Dict[str, Any]:
    results: Dict[str, Any] = {}

    steps: list[tuple[str, str, str, RecoveryRunner]] = [
        (
            "module_async",
            settings.MODULE_ASYNC_RECOVERY_LOCK_KEY,
            "module_async_recovery",
            _run_module_async_once,
        ),
        (
            "shot_split",
            settings.SHOT_SPLIT_RECOVERY_LOCK_KEY,
            "shot_split_recovery",
            _run_shot_split_once,
        ),
        (
            "generation",
            settings.GENERATION_RECOVERY_LOCK_KEY,
            "generation_recovery",
            _run_generation_once,
        ),
        (
            "download",
            settings.DOWNLOAD_RECOVERY_LOCK_KEY,
            "download_recovery",
            _run_download_once,
        ),
        (
            "video_upscale",
            settings.VIDEO_UPSCALE_RECOVERY_LOCK_KEY,
            "video_upscale_recovery",
            _run_video_upscale_once,
        ),
    ]

    for name, lock_key, log_context, runner in steps:
        try:
            results[name] = await _run_with_execution_lock(
                lock_key=lock_key,
                log_context=log_context,
                runner=runner,
            )
        except Exception as exc:
            logger.exception("启动容灾步骤执行失败。step=%s", name)
            results[name] = {"error": str(exc)}

    return {"steps": results}


if celery_app:

    @celery_app.task(
        name="recovery.startup_recovery_once",
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
    )
    def startup_recovery_once(self) -> Dict[str, Any]:
        return run_async(_run_startup_recovery_once())


    @celery_app.task(
        name="generation.recover_download_tasks_once",
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
    )
    def recover_download_tasks_once(self) -> Dict[str, Any]:
        acquired, reason = run_async(_acquire_download_recovery_loop_lock())
        if not acquired:
            return {"skipped": reason}
        try:
            result = run_async(
                _run_with_execution_lock(
                    lock_key=settings.DOWNLOAD_RECOVERY_LOCK_KEY,
                    log_context="download_recovery",
                    runner=_run_download_once,
                )
            )
            result["loop_lock"] = reason
            return result
        finally:
            _schedule_next_download_recovery_loop()


    @celery_app.task(
        name="generation.recover_generation_tasks_once",
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
    )
    def recover_generation_tasks_once(self) -> Dict[str, Any]:
        return run_async(
            _run_with_execution_lock(
                lock_key=settings.GENERATION_RECOVERY_LOCK_KEY,
                log_context="generation_recovery",
                runner=_run_generation_once,
            )
        )


    @celery_app.task(
        name="generation.dispatch_due_poll_tasks",
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
    )
    def dispatch_due_poll_tasks(self) -> Dict[str, Any]:
        return run_async(
            _run_with_execution_lock(
                lock_key=settings.POLL_DUE_DISPATCH_LOCK_KEY,
                log_context="due_poll_dispatch",
                runner=_run_due_poll_dispatch_with_guard,
                ttl_seconds=int(settings.POLL_DUE_DISPATCH_LOCK_TTL_SECONDS or 55),
            )
        )

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

    startup_recovery_once = _DisabledTask()
    recover_download_tasks_once = _DisabledTask()
    recover_generation_tasks_once = _DisabledTask()
    dispatch_due_poll_tasks = _DisabledTask()