"""
BatchScheduler - 静态批处理调度器

蓄水池模式：在时间窗口内收集请求，攒够一批后一次性调用 generate_batch，
比逐个请求串行处理更高效地利用 GPU。

典型用法:
    scheduler = BatchScheduler(backend, max_batch_size=8, max_wait_s=1.0)
    await scheduler.start()

    # 每个请求 submit 后会自动被攒批处理
    text, pt, ct = await scheduler.submit(messages, config)
"""

import asyncio
from typing import List, Optional

from loguru import logger

from .base import BaseLLMBackend, ChatMessage, GenerateConfig


class BatchScheduler:
    """蓄水池批处理调度器

    Args:
        backend: LLM 后端（需要有 generate_batch_sync 方法）
        max_batch_size: 最大批次大小
        max_wait_s: 最大等待时间（秒），收到第一个请求后最多等这么久
    """

    def __init__(
        self,
        backend: BaseLLMBackend,
        max_batch_size: int = 8,
        max_wait_s: float = 1.0,
    ):
        self._backend = backend
        self._max_batch_size = max_batch_size
        self._max_wait_s = max_wait_s
        self._queue: asyncio.Queue = None
        self._task: Optional[asyncio.Task] = None

    async def start(self):
        """启动调度器"""
        self._queue = asyncio.Queue()
        self._task = asyncio.create_task(self._worker())
        logger.info(
            f"BatchScheduler started: max_batch={self._max_batch_size}, "
            f"max_wait={self._max_wait_s}s"
        )

    async def stop(self):
        """停止调度器"""
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def submit(
        self,
        messages: List[ChatMessage],
        config: GenerateConfig,
    ) -> tuple[str, int, int]:
        """提交请求，等待批处理完成后返回结果"""
        future = asyncio.get_event_loop().create_future()
        await self._queue.put((messages, config, future))
        return await future

    async def _worker(self):
        """后台 worker：攒批 → 处理 → 分发结果"""
        while True:
            batch = []

            # 阻塞等待第一个请求
            item = await self._queue.get()
            batch.append(item)

            # 在时间窗口内继续收集
            deadline = asyncio.get_event_loop().time() + self._max_wait_s
            while len(batch) < self._max_batch_size:
                remaining = deadline - asyncio.get_event_loop().time()
                if remaining <= 0:
                    break
                try:
                    item = await asyncio.wait_for(
                        self._queue.get(), timeout=remaining
                    )
                    batch.append(item)
                except asyncio.TimeoutError:
                    break

            # 处理这一批
            if len(batch) > 1:
                logger.info(f"BatchScheduler: processing batch of {len(batch)} requests")

            await self._process_batch(batch)

    async def _process_batch(self, batch):
        """执行批处理并分发结果"""
        requests = [(msgs, cfg) for msgs, cfg, _ in batch]
        futures = [f for _, _, f in batch]

        try:
            loop = asyncio.get_event_loop()
            results = await loop.run_in_executor(
                None, self._backend.generate_batch_sync, requests
            )

            for future, result in zip(futures, results):
                if not future.done():
                    future.set_result(result)

        except Exception as e:
            logger.error(f"BatchScheduler error: {e}")
            for future in futures:
                if not future.done():
                    future.set_exception(e)
