from celery import Celery
from app.config import settings

from celery.signals import worker_process_init, worker_process_shutdown

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


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")
    celery_app.conf.update(
        broker_url=broker_url,
        result_backend=backend_url or broker_url,
        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,
        worker_prefetch_multiplier=1,
        broker_transport_options={
            "visibility_timeout": 3600,
        },
        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"},
            "app.tasks.cleanup.*": {"queue": "default"},
        },
    )
    celery_app.autodiscover_tasks(["app.tasks"])
else:
    celery_app = None


@worker_process_init.connect
def on_worker_process_init(**kwargs):
    """
    Linux prefork 子进程启动后执行。

    目的：
    1. 丢弃 fork 前可能继承的连接池状态。
    2. 后续任务会在当前子进程自己的长期 event loop 上重新建连接池。
    """
    try:
        run_async(engine.dispose())
    except Exception:
        pass


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