import logging

from celery import Celery
from celery.signals import worker_process_init, worker_process_shutdown, worker_ready

from app.config import settings
from app.models.base import engine
from app.tasks.async_runner import close_loop, run_async

logger = logging.getLogger("video_gen")


# 显式注册所有 Celery 任务模块，避免新增任务文件后 worker 启动时未注册任务。
# 不再依赖 app.tasks.__init__ 内部 import，也不再依赖 autodiscover_tasks。
CELERY_TASK_IMPORTS = (
    "app.tasks.generation_create_tasks",
    "app.tasks.generation_poll_tasks",
    "app.tasks.generation_download_tasks",
    "app.tasks.generation_recovery_tasks",
    "app.tasks.hot_opening_replicate_tasks",
    "app.tasks.shot_replicate_tasks",
    "app.tasks.shot_replicate_flow_tasks",
    "app.tasks.module_async_recovery_tasks",
    "app.tasks.user_oauth_tasks",
    "app.tasks.cleanup",
)


RECOVERY_QUEUE = settings.CELERY_RECOVERY_QUEUE or "gen_recovery"


def _derive_redis_db(url: str, db_no: int) -> str:
    if not url:
        return url
    import re

    if re.search(r"/\d+$", url):
        return re.sub(r"/\d+$", f"/{db_no}", url)
    return url.rstrip("/") + f"/{db_no}"


broker_url = settings.CELERY_BROKER_URL or (_derive_redis_db(settings.REDIS_URL, 1) if settings.REDIS_URL else "")
backend_url = settings.CELERY_RESULT_BACKEND or (_derive_redis_db(settings.REDIS_URL, 2) if settings.REDIS_URL else "")

if broker_url:
    celery_app = Celery("videogen", include=CELERY_TASK_IMPORTS)
    celery_app.conf.update(
        broker_url=broker_url,
        result_backend=backend_url or broker_url,
        imports=CELERY_TASK_IMPORTS,
        task_serializer="json",
        accept_content=["json"],
        result_serializer="json",
        timezone="Asia/Shanghai",
        enable_utc=True,
        task_soft_time_limit=600,
        task_time_limit=900,
        task_acks_late=True,
        task_reject_on_worker_lost=True,
        task_track_started=True,
        task_annotations={
            # 生成链路任务以数据库状态为准，不依赖 Celery result backend。
            # 这里忽略结果可避免任务误返回 ORM / 非 JSON 对象时触发结果序列化失败。
            # "generation.chatapi_create_generation_task": {"ignore_result": True},
            # "generation.poll_generation_task": {"ignore_result": True},
            # "generation.download_generation_result_task": {"ignore_result": True},
            "hot_opening.start_image_prompt_optimize": {"ignore_result": True},
            "hot_opening.start_video_prompt_optimize": {"ignore_result": True},
            "shot_replicate.analyze_original_video": {"ignore_result": True},
            "shot_replicate.analyze_custom_segment_video": {"ignore_result": True},
            "shot_replicate.split_one_segment": {"ignore_result": True},
            "shot_replicate.start_image_prompt_optimize": {"ignore_result": True},
            "shot_replicate.start_video_prompt_optimize": {"ignore_result": True},
        },
        worker_prefetch_multiplier=1,
        broker_transport_options={
            "visibility_timeout": 3600,
            "queue_order_strategy": "priority",
            "priority_steps": list(range(10)),
            "sep": ":",
        },
        task_routes={
            "generation.chatapi_create_generation_task": {"queue": "gen_chatapi_create"},
            "generation.poll_generation_task": {"queue": "gen_provider_poll"},
            "generation.download_generation_result_task": {"queue": "gen_result_download"},
            "hot_opening.start_image_prompt_optimize": {"queue": "gen_chatapi_create"},
            "hot_opening.start_video_prompt_optimize": {"queue": "gen_chatapi_create"},
            "shot_replicate.analyze_original_video": {"queue": "gen_chatapi_create"},
            "shot_replicate.analyze_custom_segment_video": {"queue": "gen_chatapi_create"},
            "shot_replicate.split_one_segment": {"queue": "gen_result_download"},
            "shot_replicate.start_image_prompt_optimize": {"queue": "gen_chatapi_create"},
            "shot_replicate.start_video_prompt_optimize": {"queue": "gen_chatapi_create"},
            # 恢复扫描统一走独立队列，避免占用下载/轮询/创建业务 worker。
            "recovery.startup_recovery_once": {"queue": RECOVERY_QUEUE},
            "shot_replicate.recover_split_tasks_once": {"queue": RECOVERY_QUEUE},
            "generation.recover_download_tasks_once": {"queue": RECOVERY_QUEUE},
            "generation.recover_generation_tasks_once": {"queue": RECOVERY_QUEUE},
            "module_async.recover_module_async_tasks_once": {"queue": RECOVERY_QUEUE},
            "user_oauth.update_oauth_accounts": {"queue": "default"},
            "app.tasks.cleanup.*": {"queue": "default"},
        },
    )
else:
    celery_app = None


async def _try_acquire_startup_recovery_lock() -> bool:
    """任意 worker 启动时都可尝试抢恢复投递锁，避免依赖 hostname 命名。"""
    from app.services.redis_registry_service import redis_acquire_lock

    token = await redis_acquire_lock(
        lock_key=settings.CELERY_STARTUP_RECOVERY_LOCK_KEY,
        ttl_seconds=int(settings.CELERY_STARTUP_RECOVERY_LOCK_TTL_SECONDS or 120),
        log_context="celery_startup_recovery",
    )
    return bool(token)


@worker_ready.connect
def on_worker_ready(sender=None, **kwargs):
    """Celery worker 启动时做一次容灾恢复。

    注意：
    - 不启用 Celery beat。
    - 启动容灾保留，但只投递一个 recovery.startup_recovery_once 协调任务。
    - 协调任务走独立 gen_recovery 队列，串行扫描并把真实业务任务投回原队列。
    - 所有 worker 都尝试抢 Redis 投递锁，只有抢到锁的 worker 投递恢复任务。
    """
    if celery_app is None:
        return
    if not bool(getattr(settings, "CELERY_STARTUP_RECOVERY_ENABLED", True)):
        logger.info("启动容灾恢复已关闭。CELERY_STARTUP_RECOVERY_ENABLED=false")
        return

    try:
        if not run_async(_try_acquire_startup_recovery_lock()):
            return
    except Exception:
        # Redis 不可用时不阻塞 worker 启动，避免影响稳定生成链路。
        logger.exception("启动容灾恢复锁获取失败，已跳过本次自动恢复投递")
        return

    try:
        from app.tasks.generation_recovery_tasks import startup_recovery_once

        countdown = max(0, int(settings.DOWNLOAD_RECOVERY_STARTUP_DELAY_SECONDS or 0))
        startup_recovery_once.apply_async(
            countdown=countdown,
            queue=RECOVERY_QUEUE,
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )
        logger.info(
            "启动容灾恢复协调任务已投递。queue=%s countdown=%s",
            RECOVERY_QUEUE,
            countdown,
        )
    except Exception:
        logger.exception("启动容灾恢复协调任务投递失败")


@worker_process_init.connect
def on_worker_process_init(**kwargs):
    """Linux prefork 子进程启动后丢弃 fork 前可能继承的连接池状态。"""
    try:
        run_async(engine.dispose())
    except Exception:
        pass


@worker_process_shutdown.connect
def on_worker_process_shutdown(**kwargs):
    """子进程退出前关闭连接池、Redis 注册表连接和 event loop。"""
    try:
        run_async(engine.dispose())
    except Exception:
        pass

    try:
        from app.services.redis_registry_service import close_registry_redis

        run_async(close_registry_redis())
    except Exception:
        pass
    finally:
        close_loop()