"""自适应线程池。

用于对耗时的独立子任务（如逐帧 VLM OCR）做弹性并发加速：
- 以滚动时间窗口统计已完成任务的平均响应时间；
- 窗口内平均响应 < fast_threshold（默认 0.3s）→ 增加 1 个工作线程（上限 max_workers）；
- 窗口内平均响应 > slow_threshold（默认 1.0s）→ 减少 1 个工作线程（下限 min_workers）。

线程数从 min_workers（默认 1）起步，按实测负载自适应：服务端空闲（响应快）
就加大并发，服务端变慢就退避，避免盲目并发压垮上游（如本地 Ollama）。

线程安全说明：worker 会在多个线程中并发调用，调用方需保证 worker 无共享
可变状态（registry 处理器是纯函数，符合要求）；结果按输入顺序返回。
"""

from __future__ import annotations

import queue
import threading
import time
from typing import Callable

# 停止哨兵：压入队列让空闲工作线程退出（用于缩容）。
_POISON = object()


def decide(
    current: int,
    avg: float,
    min_workers: int,
    max_workers: int,
    fast_threshold: float,
    slow_threshold: float,
) -> int:
    """根据窗口平均响应时间返回调整后的目标线程数（纯决策函数）。

    响应快（avg < fast_threshold）且未达上限 → 加 1；响应慢
    （avg > slow_threshold）且未达下限 → 减 1；其余情况保持不变。
    """
    if avg < fast_threshold and current < max_workers:
        return current + 1
    if avg > slow_threshold and current > min_workers:
        return current - 1
    return current


class AdaptiveThreadPool:
    """自适应线程池：单次 map 按输入顺序返回全部结果。"""

    def __init__(
        self,
        worker: Callable,
        min_workers: int = 1,
        max_workers: int = 16,
        window_seconds: float = 10.0,
        fast_threshold: float = 0.3,
        slow_threshold: float = 1.0,
        clock=time.monotonic,
        on_progress: Callable[[int, int, float], None] | None = None,
    ) -> None:
        """初始化；clock 可注入便于测试；on_progress(done,total,rate) 每次完成回调。"""
        self._worker = worker
        self.min_workers = max(1, min_workers)
        self.max_workers = max(self.min_workers, max_workers)
        self.window_seconds = window_seconds
        self.fast_threshold = fast_threshold
        self.slow_threshold = slow_threshold
        self._clock = clock
        self._queue: queue.Queue = queue.Queue()
        # 并发目标线程数：决策/缩容的权威依据（线程退出是异步的，不能用
        # len(_threads) 判断，否则并发缩容会重复放哨兵把全部线程毒死）。
        self._target_workers = 0
        self._threads: list[threading.Thread] = []
        self._results: list = []
        self._lock = threading.Lock()
        self._stop = threading.Event()
        # 滚动窗口起点与已记录的单次耗时。
        self._window_start = clock()
        # 观测到的最大并发线程数（供测试与监控）。
        self.max_concurrency = 0
        # 进度回调与计数：on_progress(已完成数, 总数, 平均速度/秒)。
        self._on_progress = on_progress
        self._completed = 0
        self._total = 0
        self._started_at = 0.0
        self._window_times: list[float] = []

    def _run(self) -> None:
        """工作线程主循环：取任务 → 执行 → 记录耗时并自适应评估。"""
        try:
            while not self._stop.is_set():
                try:
                    seq, item = self._queue.get(timeout=0.2)
                except queue.Empty:
                    continue
                if item is _POISON:
                    # 缩容哨兵：处理完即可退出（队列计数照常）。
                    self._queue.task_done()
                    break
                start = self._clock()
                try:
                    result = self._worker(item)
                except Exception as exc:
                    # 单任务异常不拖垮整体：以异常对象作为结果，由调用方判定。
                    result = exc
                finally:
                    elapsed = self._clock() - start
                self._results.append((seq, result))
                # 进度回调：已完成数、总数与平均处理速度（条/秒）。
                self._completed += 1
                if self._on_progress is not None:
                    elapsed_total = max(self._clock() - self._started_at, 1e-9)
                    self._on_progress(
                        self._completed, self._total, self._completed / elapsed_total
                    )
                self._tick(elapsed)
                self._queue.task_done()
        finally:
            # 无论何种退出路径都从线程列表移除，保证线程数统计准确。
            with self._lock:
                if threading.current_thread() in self._threads:
                    self._threads.remove(threading.current_thread())

    def _tick(self, elapsed: float) -> None:
        """记录一次完成耗时；窗口满时按平均响应时间调整线程数。"""
        self._window_times.append(elapsed)
        if self._clock() - self._window_start < self.window_seconds:
            return
        avg = sum(self._window_times) / len(self._window_times)
        self._window_start = self._clock()
        self._window_times.clear()
        with self._lock:
            current = self._target_workers
        self._resize(
            decide(
                current, avg, self.min_workers, self.max_workers,
                self.fast_threshold, self.slow_threshold,
            )
        )

    def _resize(self, target: int) -> None:
        """调整并发目标：扩容启动新线程；缩容压入等量停止哨兵（幂等）。

        以 _target_workers 为当前值：重复调用同一 target 不会重复放哨兵，
        避免并发缩容把所有线程毒死导致队列任务无人处理而挂起。
        """
        with self._lock:
            current = self._target_workers
            if target > current:
                self.max_concurrency = max(self.max_concurrency, target)
                for _ in range(target - current):
                    thread = threading.Thread(target=self._run, daemon=True)
                    thread.start()
                    self._threads.append(thread)
                self._target_workers = target
            elif target < current:
                for _ in range(current - target):
                    self._queue.put((None, _POISON))
                self._target_workers = target

    def map(self, items) -> list:
        """按输入顺序返回每个 item 经 worker 处理后的结果列表。"""
        self._results = []
        self._completed = 0
        self._total = len(items)
        self._started_at = self._clock()
        self._stop.clear()
        self._resize(self.min_workers)
        for seq, item in enumerate(items):
            self._queue.put((seq, item))
        self._queue.join()
        self._stop.set()
        with self._lock:
            threads = list(self._threads)
        for thread in threads:
            thread.join(1.0)
        self._results.sort(key=lambda pair: pair[0])
        return [result for _, result in self._results]