from loguru import logger
from typing import List, Iterator
import asyncio
from .base_processor import BaseProcessor
from .pipeline_data import PipelineData


class Pipeline:
    """管道核心类，管理多个处理器，负责数据流转"""

    def __init__(self):
        self.processors: List[BaseProcessor] = []  # 同步处理器列表
        self.async_processors: List[BaseProcessor] = []  # 异步处理器列表
        # 初始化日志记录器
        self.logger = logger
        self.logger.debug("管道已初始化")

    def add_processor(self, processor: BaseProcessor) -> None:
        """向管道添加同步处理器（按添加顺序执行）"""
        self.processors.append(processor)
        self.logger.info(f"管道已添加同步处理器: {processor}")

    def add_processor_async(self, processor: BaseProcessor) -> None:
        """向管道添加异步处理器（异步执行，不阻塞数据流）"""
        self.async_processors.append(processor)
        self.logger.info(f"管道已添加异步处理器: {processor}")

    def remove_processor(self, processor_name: str) -> bool:
        """移除指定名称的处理器"""
        # 检查同步处理器列表
        for idx, p in enumerate(self.processors):
            if p.name == processor_name:
                self.processors.pop(idx)
                self.logger.info(f"管道已移除同步处理器: {p}")
                return True
        # 检查异步处理器列表
        for idx, p in enumerate(self.async_processors):
            if p.name == processor_name:
                self.async_processors.pop(idx)
                self.logger.info(f"管道已移除异步处理器: {p}")
                return True
        self.logger.warning(f"未找到处理器: {processor_name}")
        return False

    def run(self, data_iterator: Iterator[PipelineData]) -> Iterator[PipelineData]:
        """
        运行管道：迭代输入数据，依次经过所有同步处理器，同时异步执行异步处理器
        :param data_iterator: 数据源迭代器（如视频+YOLO的迭代器）
        :return: 处理后的数据包迭代器
        """

        # 创建异步任务列表，用于管理异步处理器
        async def run_async_processors(data: PipelineData):
            tasks = []
            for async_processor in self.async_processors:
                task = asyncio.create_task(self._process_async(async_processor, data))
                tasks.append(task)
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)

        for data in data_iterator:
            # 依次执行每个同步处理器的处理逻辑
            for processor in self.processors:
                try:
                    data = processor.process(data)
                except Exception as e:
                    self.logger.error(f"同步处理器 {processor.name} 执行出错: {e}")

            # 异步执行异步处理器，不阻塞数据流
            if self.async_processors:
                # 在独立的事件循环中运行异步处理器（如果可用）
                try:
                    # 尝试获取当前事件循环
                    loop = asyncio.get_running_loop()
                    # 在当前事件循环中调度异步处理器
                    asyncio.create_task(run_async_processors(data.copy() if hasattr(data, 'copy') else data))
                except RuntimeError:
                    # 如果没有运行中的事件循环，则创建一个新的
                    asyncio.run(run_async_processors(data.copy() if hasattr(data, 'copy') else data))

            yield data  # 返回处理后的数据包

    async def _process_async(self, processor, data):
        """异步执行单个处理器"""
        try:
            # 对数据进行异步处理（如果处理器支持）
            if hasattr(processor, 'process_async'):
                await processor.process_async(data)
            else:
                # 如果处理器不支持异步处理，则同步处理
                processor.process(data)
        except Exception as e:
            self.logger.error(f"异步处理器 {processor.name} 执行出错: {e}")