import asyncio
import inspect
import threading
import time
from collections.abc import Iterator
from concurrent.futures import Executor, Future, ThreadPoolExecutor, wait
from contextlib import contextmanager
from contextvars import Context, copy_context
from functools import partial, wraps
from typing import Any, Awaitable, Callable, ParamSpec, TypeVar, cast

from anyio.to_thread import run_sync

from app.runtime.gc import flush_thread_allocator_cache
from app.schemas.exception import ImmediateException

TaskResult = TypeVar("TaskResult")
ExecutorResult = TypeVar("ExecutorResult")
_Params = ParamSpec("_Params")
_Result = TypeVar("_Result")

# 包目录替换在线程池中执行，导入在同步生命周期线程中执行；两者须共享进程锁。
_plugin_package_locks: dict[str, Any] = {}
_plugin_package_locks_guard = threading.Lock()


@contextmanager
def plugin_package_lock(plugin_id: str) -> Iterator[None]:
    """串行化单个插件的代码导入、包替换和目录恢复。"""
    normalized_id = plugin_id.casefold()
    with _plugin_package_locks_guard:
        lock = _plugin_package_locks.setdefault(
            normalized_id,
            threading.RLock(),
        )
    with lock:
        yield


def submit_with_context(
        executor: Executor,
        func: Callable[..., ExecutorResult],
        /,
        *args: Any,
        **kwargs: Any,
) -> Future[ExecutorResult]:
    """从空线程上下文提交任务，并在执行时恢复调用方的独立上下文快照。"""
    context = copy_context()

    def submit() -> Future[ExecutorResult]:
        """在无继承上下文中提交，再由 worker 执行已捕获的调用方快照。"""
        return executor.submit(context.run, func, *args, **kwargs)

    return Context().run(submit)


class IdleCacheReleasingThreadPoolExecutor(ThreadPoolExecutor):
    """
    worker 执行完任务、且队列里暂无待办时，由该 worker 交还自己的分配器线程缓存。

    标准库线程池只按需增长、从不收缩，高峰期建出的 worker 此后长期空闲；jemalloc 只在线程
    继续分配时才回收其线程缓存，空闲 worker 在最后一个任务里囤积的空闲块会一直计入 RSS。
    刷新只发生在即将转入空闲时：连续任务之间不刷新，不影响突发期的分配性能。
    """

    def submit(
            self,
            fn: Callable[..., ExecutorResult],
            /,
            *args: Any,
            **kwargs: Any,
    ) -> Future[ExecutorResult]:
        """提交任务，任务结束后在同一 worker 线程里按需交还线程缓存。"""
        return super().submit(self._run_then_release_cache, fn, *args, **kwargs)

    def _run_then_release_cache(
            self,
            fn: Callable[..., ExecutorResult],
            /,
            *args: Any,
            **kwargs: Any,
    ) -> ExecutorResult:
        """在 worker 线程中执行任务；队列已空说明该 worker 即将等待，此时刷新其线程缓存。"""
        try:
            return fn(*args, **kwargs)
        finally:
            # _work_queue 是标准库线程池的待办队列，自 Python 3.2 起稳定存在
            if self._work_queue.empty():
                flush_thread_allocator_cache()


class OwnedThreadPoolExecutor(IdleCacheReleasingThreadPoolExecutor):
    """
    追踪已接受 Future，并提供可重试的有界关闭合同。

    该 owner 不取消排队任务，保持 ``ThreadPoolExecutor.shutdown(wait=True)``
    的历史完成语义；区别仅在于调用方可以在预算耗尽后保留同一实例继续收敛。
    """

    def __init__(self, max_workers: int | None = None) -> None:
        """初始化线程池、提交准入状态和 Future owner 集合。"""
        super().__init__(max_workers=max_workers)
        self._ownership_lock = threading.RLock()
        self._accepting = True
        self._owned_futures: set[Future[Any]] = set()

    @property
    def accepting(self) -> bool:
        """返回执行器是否仍允许提交新任务。"""
        with self._ownership_lock:
            return self._accepting

    def submit(
            self,
            fn: Callable[..., ExecutorResult],
            /,
            *args: Any,
            **kwargs: Any,
    ) -> Future[ExecutorResult]:
        """提交任务并在其达到终态前保留 owner。"""
        with self._ownership_lock:
            future = super().submit(fn, *args, **kwargs)
            self._owned_futures.add(future)
            future.add_done_callback(self._discard_future)
            return future

    def shutdown(
            self,
            wait: bool = True,
            *,
            cancel_futures: bool = False,
    ) -> None:
        """封口提交准入并保留标准库 shutdown 的调用语义。"""
        with self._ownership_lock:
            self._accepting = False
            # 先在锁内封口；真正等待必须在锁外进行，否则 worker 的完成回调无法释放 owner。
            super().shutdown(wait=False, cancel_futures=cancel_futures)
        if wait:
            super().shutdown(wait=True, cancel_futures=cancel_futures)

    def _discard_future(self, future: Future[Any]) -> None:
        """任务达到终态后释放 owner 记录。"""
        with self._ownership_lock:
            self._owned_futures.discard(future)

    def shutdown_bounded(self, timeout: float) -> bool:
        """
        封口新提交并有限等待全部已接受任务。

        :param timeout: 等待 Future 达到终态的最长秒数
        :return: 所有任务与 worker 均已终止时返回 True，否则返回 False
        """
        deadline = time.monotonic() + max(0.0, timeout)
        with self._ownership_lock:
            self._accepting = False
            # 不取消排队工作，保持历史 shutdown(wait=True) 的完成语义。
            super().shutdown(wait=False)
            owned_futures = tuple(self._owned_futures)
        if owned_futures:
            _, pending_futures = wait(
                owned_futures,
                timeout=max(0.0, deadline - time.monotonic()),
            )
            if pending_futures:
                return False
        # 标准库只提供无界 wait=True；Future 又会先标记完成再执行 done callback，
        # 因此封口后读取稳定 worker 集合，复用同一 deadline 做有限 join。
        worker_threads = tuple(cast(set[threading.Thread], self._threads))
        current_thread = threading.current_thread()
        for worker_thread in worker_threads:
            if worker_thread is current_thread:
                continue
            worker_thread.join(
                timeout=max(0.0, deadline - time.monotonic()),
            )
        return all(not worker_thread.is_alive() for worker_thread in worker_threads)


async def await_task_to_terminal(
    task: asyncio.Future[TaskResult],
) -> TaskResult:
    """忽略当前调用方的重复取消，直到受保护任务进入真实终态。"""
    while not task.done():
        try:
            await asyncio.shield(task)
        except asyncio.CancelledError:
            continue
        except BaseException:
            break
    return task.result()


async def run_in_threadpool(
    func: Callable[..., Any],
    *args: Any,
    **kwargs: Any,
) -> Any:
    """在线程中执行同步函数，保持 FastAPI 旧帮助函数的参数语义。"""
    if kwargs:
        func = partial(func, **kwargs)
    context = copy_context()
    return await run_sync(context.run, func, *args)


async def run_in_threadpool_to_completion(
    func: Callable[..., Any],
    *args: Any,
    **kwargs: Any,
) -> Any:
    """在线程调用取得终态后传播取消，避免提前释放仍在使用的执行容量。"""
    worker_task = asyncio.create_task(run_in_threadpool(func, *args, **kwargs))
    cancellation: asyncio.CancelledError | None = None
    while not worker_task.done():
        try:
            await asyncio.wait({worker_task})
        except asyncio.CancelledError as error:
            cancellation = cancellation or error
            continue
    try:
        result = worker_task.result()
    except Exception as error:
        if cancellation is not None:
            raise cancellation from error
        raise
    if cancellation is not None:
        raise cancellation
    return result


def retry(
    ExceptionToCheck: type[BaseException] | tuple[type[BaseException], ...],
    tries: int = 3,
    delay: int = 3,
    backoff: int = 2,
    logger: Any = None,
) -> Callable[[Callable[_Params, _Result]], Callable[_Params, _Result]]:
    """
    :param ExceptionToCheck: 需要捕获的异常
    :param tries: 重试次数
    :param delay: 延迟时间
    :param backoff: 延迟倍数
    :param logger: 日志对象
    """

    def deco_retry(f: Callable[_Params, _Result]) -> Callable[_Params, _Result]:
        """保留调用参数与返回类型，按原有检测方式选择同步或异步重试。"""
        def f_retry(*args: _Params.args, **kwargs: _Params.kwargs) -> _Result:
            """按既有退避间隔重试可捕获异常，立即终止异常直接交给调用方。"""
            mtries, mdelay = tries, delay
            while mtries > 1:
                try:
                    return f(*args, **kwargs)
                except ImmediateException:
                    raise
                except ExceptionToCheck as e:
                    msg = f"{str(e)}, {mdelay} 秒后重试 ..."
                    if logger:
                        logger.warn(msg)
                    else:
                        print(msg)
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return f(*args, **kwargs)

        async def async_f_retry(*args: _Params.args, **kwargs: _Params.kwargs) -> object:
            """仅在原函数被识别为协程时调用，异步等待与同步分支相同的重试策略。"""
            mtries, mdelay = tries, delay
            while mtries > 1:
                try:
                    return await cast(Awaitable[object], f(*args, **kwargs))
                except ImmediateException:
                    raise
                except ExceptionToCheck as e:
                    msg = f"{str(e)}, {mdelay} 秒后重试 ..."
                    if logger:
                        logger.warn(msg)
                    else:
                        print(msg)
                    await asyncio.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return await cast(Awaitable[object], f(*args, **kwargs))

        # 根据函数类型返回相应的包装器
        if inspect.iscoroutinefunction(f):
            # 仅协程分支可调用异步包装器，原 await 结果与参数原样透传。
            return cast(Callable[_Params, _Result], async_f_retry)
        else:
            return f_retry

    return deco_retry


def log_execution_time(
    logger: Any = None,
) -> Callable[[Callable[_Params, _Result]], Callable[_Params, _Result]]:
    """
    记录函数执行时间的装饰器
    :param logger: 日志记录器对象，用于记录异常信息
    """

    def decorator(func: Callable[_Params, _Result]) -> Callable[_Params, _Result]:
        """保留原函数元数据及调用合同，按原分派方式记录执行耗时。"""
        @wraps(func)
        def wrapper(*args: _Params.args, **kwargs: _Params.kwargs) -> _Result:
            """同步函数成功后记录耗时，结果和异常沿用原传播路径。"""
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            msg = f"{func.__name__} execution time: {end_time - start_time:.2f} seconds"
            if logger:
                logger.debug(msg)
            else:
                print(msg)
            return result

        @wraps(func)
        async def async_wrapper(*args: _Params.args, **kwargs: _Params.kwargs) -> object:
            """协程成功结束后记录耗时，异常不产生成功耗时日志。"""
            start_time = time.time()
            result = await cast(Awaitable[object], func(*args, **kwargs))
            end_time = time.time()
            msg = f"{func.__name__} execution time: {end_time - start_time:.2f} seconds"
            if logger:
                logger.debug(msg)
            else:
                print(msg)
            return result

        # 根据函数类型返回相应的包装器
        if inspect.iscoroutinefunction(func):
            # 已确认原函数返回协程，异步包装器保留同一 await 结果。
            return cast(Callable[_Params, _Result], async_wrapper)
        else:
            return wrapper

    return decorator