import asyncio
from collections import defaultdict
from typing import Callable, Coroutine, Dict, Any, Type
from loguru import logger

from src.tasks.surugaya import SurugayaPurchaseTask
from src.tasks.surugaya_trade_monitor import SurugayaTradeMonitorTask

# topic 到采购任务类的映射注册表
TOPIC_TO_TASK_MAP: Dict[str, Type[Any]] = {
    "surugaya": SurugayaPurchaseTask,  # 骏河屋采购任务
    "surugaya_trade_monitor": SurugayaTradeMonitorTask,  # 骏河屋交易订单获取任务
}


class TaskDispatcher:
    """
    任务分发器：处理来自消息队列的采购任务。
    - 在后台协程中执行任务，与 Redis 监听器解耦。
    - 按账号维度使用 asyncio.Lock 串行化浏览器操作，避免冲突。
    - 支持按 topic 模块化路由到不同工作流。
    """

    def __init__(self) -> None:
        # 按 account_id 维护锁字典，新账号会自动创建锁
        # 单线程事件循环下完全异步安全
        self._account_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
        logger.info("任务分发器已初始化（支持按账号串行锁）。")

    async def dispatch(
            self,
            message_data: Dict[str, Any],
            on_complete_callback: Callable[[str, bool, Any], Coroutine[Any, Any, None]] = None
    ) -> None:
        """
        启动非阻塞后台协程处理任务，使监听器无需等待浏览器自动化完成即可继续轮询 Redis。
        
        Args:
            message_data: 已解析的任务消息字典。
            on_complete_callback: 任务完成（成功或失败）时执行的异步回调。
        """
        task_id = message_data.get("task_id", "N/A")
        # logger.debug(f"正在为任务 '{task_id}' 启动后台协程...")

        # 投递到事件循环，fire-and-forget
        asyncio.create_task(self._execute_task_safe(message_data, on_complete_callback))

    async def _execute_task_safe(
            self,
            message_data: Dict[str, Any],
            on_complete_callback: Callable[[str, bool, Any], Coroutine[Any, Any, None]] = None
    ) -> None:
        """
        在账号锁约束下执行具体采购流程，并统一记录异常。
        """
        task_id = message_data.get("task_id", "N/A")
        topic = str(message_data.get("topic"))
        account_id = str(message_data.get("purchase_account"))
        callback_id = message_data.get("_stream_msg_id") or message_data.get("_callback_id") or task_id

        logger.info(f"[任务 {task_id}] 开始分发，topic='{topic}'，账号='{account_id}'。")

        # 1. 校验必填字段
        if not topic or not account_id:
            logger.error(f"[任务 {task_id}] 缺少必填字段 topic 或 purchase_account，任务已中止。")
            if on_complete_callback:
                await on_complete_callback(str(callback_id), False, {"error": "缺少关键路由字段。", "result": None})
            return

        # 2. 根据 topic 匹配已注册的工作流
        task_class = TOPIC_TO_TASK_MAP.get(topic.lower())
        if not task_class:
            logger.error(f"[任务 {task_id}] 未找到 topic '{topic}' 对应的采购任务类，任务已中止。")
            if on_complete_callback:
                await on_complete_callback(
                    str(callback_id),
                    False,
                    {"error": f"未知的 topic: {topic}", "result": None},
                )
            return

        # 3. 获取该账号的串行锁
        lock = self._account_locks[account_id]
        logger.info(f"[任务 {task_id}] 账号 '{account_id}' 忙碌中，已进入串行锁队列等待...")

        # 异步锁保证同一 account_id 同一时刻只有一个协程在执行
        async with lock:
            logger.success(f"[任务 {task_id}] 已获取账号 '{account_id}' 的锁，正在启动浏览器工作流...")

            success = False
            error_msg = None
            result_data: Any = None
            task_instance: Any = None

            try:
                # 实例化任务并执行
                task_instance = task_class(message_data)
                await task_instance.run()
                success = True
                result_data = getattr(task_instance, "task_result_data", None)
                logger.success(f"[任务 {task_id}] 工作流执行成功。")
            except Exception as e:
                error_msg = str(e)
                if task_instance:
                    result_data = getattr(task_instance, "task_result_data", None)
                logger.error(f"[任务 {task_id}] 工作流因未处理异常而终止: {e}")
            finally:
                # 4. 触发完成回调（例如 Redis Stream ACK）
                if on_complete_callback:
                    try:
                        await on_complete_callback(
                            str(callback_id),
                            success,
                            {"error": error_msg, "result": result_data},
                        )
                    except Exception as cb_err:
                        logger.error(f"[任务 {task_id}] 完成回调执行失败: {cb_err}")

                logger.debug(f"[任务 {task_id}] 已释放账号 '{account_id}' 的锁。")