"""
线程执行上下文与取消边界，供运行时统一判断可终止的工作。
"""

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.runtime.priority import background_thread_limiter, interactive_thread_limiter, is_interactive
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[..., _Result],
    *args: Any,
    **kwargs: Any,
) -> _Result:
    """
    在线程中执行同步函数，保持 FastAPI 旧帮助函数的参数语义。

    交互调用使用保留配额；普通宿主调用使用独立后台配额，避免耗尽框架同步入口的令牌。

    :param func: 在线程中执行的同步回调
    :param args: 回调的位置参数
    :param kwargs: 回调的关键字参数
    :return: 回调完成后的原始结果
    """
    if kwargs:
        func = partial(func, **kwargs)
    context = copy_context()

    def invoke() -> _Result:
        """
        在线程内恢复提交方上下文，确保嵌套调用沿用执行归属。
        """
        return context.run(func, *args)

    limiter = interactive_thread_limiter() if is_interactive() else background_thread_limiter()
    result: _Result = await run_sync(invoke, limiter=limiter)
    return result


async def run_in_threadpool_to_completion(
    func: Callable[..., _Result],
    *args: Any,
    **kwargs: Any,
) -> _Result:
    """
    在线程调用取得终态后传播取消，避免提前释放仍在使用的执行容量。

    :param func: 在线程中执行的同步回调
    :param args: 回调的位置参数
    :param kwargs: 回调的关键字参数
    :return: 回调完成后的原始结果；调用方取消时先等回调结束再传播取消
    """
    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: 允许重试的异常类型或类型元组；ImmediateException 始终直接传播
    :param tries: 总尝试次数，包含首次调用；小于等于 1 时仍调用一次
    :param delay: 首次失败后的等待秒数，默认 3 秒
    :param backoff: 每次重试后等待时间的乘数，默认 2
    :param logger: 可选日志对象；未提供时将重试提示写入标准输出
    :return: 保留原调用参数和结果语义的装饰器
    """

    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