"""Agent 并发调度器"""

import logging
from concurrent.futures import ThreadPoolExecutor, Future
from typing import List

from .agent import SubAgent
from .config import SubAgentConfig

logger = logging.getLogger(__name__)


class AgentScheduler:
    """Agent 并发调度器

    负责调度多个子 Agent 的执行，支持串行和并行模式
    """

    def __init__(self, max_workers: int = 4):
        """初始化调度器

        Args:
            max_workers: 最大并行工作线程数
        """
        self._executor = ThreadPoolExecutor(max_workers=max_workers)

    def schedule_sequential(
        self,
        configs: List[SubAgentConfig],
        inputs: List[str],
    ) -> List[str]:
        """
        串行调度多个子 Agent

        Args:
            configs: 子 Agent 配置列表
            inputs: 输入列表

        Returns:
            结果列表
        """
        if len(configs) != len(inputs):
            raise ValueError("配置数量和输入数量必须一致")

        results = []
        for config, input_text in zip(configs, inputs):
            agent = SubAgent(config, depth=1)
            result = agent.invoke(input_text)
            results.append(result)

        return results

    def schedule_parallel(
        self,
        configs: List[SubAgentConfig],
        inputs: List[str],
    ) -> List[str]:
        """
        并行调度多个子 Agent

        Args:
            configs: 子 Agent 配置列表
            inputs: 输入列表

        Returns:
            结果列表（顺序与输入一致）
        """
        if len(configs) != len(inputs):
            raise ValueError("配置数量和输入数量必须一致")

        def run_agent(config: SubAgentConfig, input_text: str) -> str:
            agent = SubAgent(config, depth=1)
            return agent.invoke(input_text)

        # 使用线程池并行执行
        futures: List[Future] = [
            self._executor.submit(run_agent, config, input_text)
            for config, input_text in zip(configs, inputs)
        ]

        results = [f.result() for f in futures]
        return results

    def shutdown(self, wait: bool = True):
        """关闭调度器

        Args:
            wait: 是否等待所有任务完成
        """
        self._executor.shutdown(wait=wait)

    def __del__(self):
        """清理资源"""
        if hasattr(self, "_executor"):
            self._executor.shutdown(wait=False)
