from __future__ import annotations

import asyncio
import logging
import os
import threading
from concurrent.futures import Future, TimeoutError as FutureTimeoutError
from typing import Awaitable, TypeVar

from app.config import settings

logger = logging.getLogger("video_gen")

T = TypeVar("T")

_thread_local = threading.local()
_single_loop_lock = threading.RLock()
_single_loop: asyncio.AbstractEventLoop | None = None
_single_loop_thread: threading.Thread | None = None
_single_loop_pid: int | None = None
_single_loop_ready: threading.Event | None = None


async def _dispose_async_resources() -> None:
    """释放当前 async loop 内缓存的异步资源。

    Celery soft time limit 会打断同步等待 future.result() 的线程；如果不主动
    cancel coroutine 并释放 engine/redis，后台 loop 里残留的协程可能继续占用
    SQLAlchemy QueuePool 连接，后续任务就会出现 QueuePool timeout。
    """
    try:
        from app.services.redis_registry_service import close_registry_redis

        await close_registry_redis()
    except Exception:
        logger.debug("关闭 Celery Redis registry 连接失败", exc_info=True)

    try:
        from app.models.base import engine

        await engine.dispose()
    except Exception:
        logger.debug("dispose Celery SQLAlchemy engine 失败", exc_info=True)


def _runner_mode() -> str:
    mode = str(getattr(settings, "CELERY_ASYNC_RUNNER_MODE", "single_loop") or "single_loop").strip().lower()
    if mode not in {"single_loop", "direct"}:
        return "single_loop"
    return mode


def _get_or_create_thread_local_loop() -> asyncio.AbstractEventLoop:
    """兼容旧方案：当前线程持有一个长期 event loop。

    仅作为降级模式使用。长期推荐 single_loop，避免 Windows threads 下
    多线程 event loop 复用 asyncpg / redis.asyncio 连接对象。
    """
    pid = os.getpid()
    loop = getattr(_thread_local, "loop", None)
    loop_pid = getattr(_thread_local, "pid", None)

    if loop is None or loop.is_closed() or loop_pid != pid:
        loop = asyncio.new_event_loop()
        _thread_local.loop = loop
        _thread_local.pid = pid

    return loop


def _single_loop_worker(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None:
    asyncio.set_event_loop(loop)
    ready.set()
    try:
        loop.run_forever()
    finally:
        pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
        if pending:
            for task in pending:
                task.cancel()
            loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))

        loop.run_until_complete(_dispose_async_resources())
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.close()


def _get_or_create_single_loop() -> asyncio.AbstractEventLoop:
    """获取当前 Celery 进程内唯一 async event loop。

    Linux prefork：每个 Celery 子进程各自一个 loop。
    Windows threads：同一 worker 进程内所有任务线程共享同一个 loop。
    """
    global _single_loop, _single_loop_thread, _single_loop_pid, _single_loop_ready

    pid = os.getpid()
    with _single_loop_lock:
        if (
            _single_loop is not None
            and not _single_loop.is_closed()
            and _single_loop_thread is not None
            and _single_loop_thread.is_alive()
            and _single_loop_pid == pid
        ):
            return _single_loop

        # fork 后 pid 变化，必须丢弃父进程状态，重新创建子进程自己的 loop。
        _single_loop = asyncio.new_event_loop()
        _single_loop_pid = pid
        _single_loop_ready = threading.Event()
        _single_loop_thread = threading.Thread(
            target=_single_loop_worker,
            args=(_single_loop, _single_loop_ready),
            name=f"celery-async-runner-{pid}",
            daemon=True,
        )
        _single_loop_thread.start()
        _single_loop_ready.wait(timeout=5)
        return _single_loop


def _cancel_future_and_reset_loop(future: Future[T] | None, *, reason: str) -> None:
    """取消当前协程并重置当前进程内 event loop。"""
    if future is not None and not future.done():
        future.cancel()
        try:
            future.result(timeout=2)
        except Exception:
            pass

    logger.warning("Celery async_runner 正在重置 event loop。reason=%s", reason)
    close_loop()


def run_async(coro: Awaitable[T]) -> T:
    """Celery 同步 task 调用异步协程的统一入口。

    默认 single_loop 模式：
    - 一个 Celery 子进程只有一个专用 event loop；
    - 所有 asyncpg / redis.asyncio 操作都在这个 loop 内创建和使用；
    - 避免 got Future attached to a different loop；
    - 当 Celery soft time limit 打断 future.result() 时，主动 cancel 后台协程并
      释放连接池，避免 QueuePool 被残留任务长期占用。

    降级 direct 模式：
    - 兼容旧的线程本地 loop 方案；
    - 如果使用 direct，建议同时开启 CELERY_DB_USE_NULLPOOL=true。
    """
    if _runner_mode() == "direct":
        loop = _get_or_create_thread_local_loop()
        try:
            return loop.run_until_complete(coro)
        except BaseException:
            try:
                if not loop.is_closed():
                    loop.run_until_complete(_dispose_async_resources())
            finally:
                close_loop()
            raise

    loop = _get_or_create_single_loop()
    try:
        running_loop = asyncio.get_running_loop()
    except RuntimeError:
        running_loop = None

    if running_loop is loop:
        raise RuntimeError("run_async() 不能在 Celery async_runner 的事件循环内部被同步调用")

    future: Future[T] = asyncio.run_coroutine_threadsafe(coro, loop)
    try:
        return future.result()
    except FutureTimeoutError:
        _cancel_future_and_reset_loop(future, reason="future_result_timeout")
        raise
    except BaseException:
        # Celery SoftTimeLimitExceeded/worker shutdown 等异常会从这里抛出。
        # 必须重置 loop，否则后台协程继续运行会拖住 DB 连接池。
        _cancel_future_and_reset_loop(future, reason="base_exception")
        raise


def close_loop() -> None:
    """关闭当前进程内 async runner loop。"""
    global _single_loop, _single_loop_thread, _single_loop_pid, _single_loop_ready

    # 关闭 single_loop。
    with _single_loop_lock:
        loop = _single_loop
        thread = _single_loop_thread
        if loop is not None and not loop.is_closed() and thread is not None and thread.is_alive():
            try:
                cleanup_future = asyncio.run_coroutine_threadsafe(_dispose_async_resources(), loop)
                cleanup_future.result(timeout=5)
            except Exception:
                logger.debug("关闭 loop 前清理 async 资源失败", exc_info=True)
            try:
                loop.call_soon_threadsafe(loop.stop)
                thread.join(timeout=5)
            except Exception:
                logger.debug("关闭 Celery async_runner loop 失败", exc_info=True)

        _single_loop = None
        _single_loop_thread = None
        _single_loop_pid = None
        _single_loop_ready = None

    # 关闭 direct 降级模式的线程本地 loop。
    loop = getattr(_thread_local, "loop", None)
    if loop is not None and not loop.is_closed():
        try:
            pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
            for task in pending:
                task.cancel()
            if pending:
                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
            loop.run_until_complete(_dispose_async_resources())
            loop.run_until_complete(loop.shutdown_asyncgens())
        except Exception:
            logger.debug("关闭 direct loop 前清理失败", exc_info=True)
        finally:
            loop.close()
    _thread_local.loop = None
    _thread_local.pid = None
