#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
优先级队列管理模块
实现动态阈值多级队列和平衡调度策略
"""

import asyncio
import logging
import heapq
from typing import Dict, List, Any, Optional, Tuple, Set
from datetime import datetime, timedelta
from collections import defaultdict
from utils.logger import get_logger

# 获取日志记录器
logger = get_logger("services.enhanced_scheduler.priority_queue_manager")

class PriorityQueueManager:
    """
    优先级队列管理器
    实现动态阈值多级队列和平衡调度策略
    """
    
    def __init__(self, 
                 queue_count: int = 3, 
                 threshold_percentiles: List[float] = None,
                 worker_ratios: List[float] = None,
                 rebalance_interval: int = 300):
        """
        初始化优先级队列管理器
        
        Args:
            queue_count: 队列数量
            threshold_percentiles: 队列阈值百分比，如[0.1, 0.3, 1.0]表示前10%、接下来20%、剩余的70%
            worker_ratios: 工作线程分配比例，如[0.6, 0.3, 0.1]
            rebalance_interval: 重新平衡间隔（秒）
        """
        self.queue_count = queue_count
        self.threshold_percentiles = threshold_percentiles
        self.worker_ratios = worker_ratios
        self.rebalance_interval = rebalance_interval
        
        # 验证配置
        if len(self.threshold_percentiles) != queue_count:
            raise ValueError(f"阈值百分比数量({len(self.threshold_percentiles)})必须等于队列数量({queue_count})")
        if len(self.worker_ratios) != queue_count:
            raise ValueError(f"工作线程比例数量({len(self.worker_ratios)})必须等于队列数量({queue_count})")
        if sum(self.worker_ratios) != 1.0:
            raise ValueError(f"工作线程比例总和必须等于1.0，当前为{sum(self.worker_ratios)}")
        
        # 初始化多级队列
        self.priority_queues = [asyncio.PriorityQueue() for _ in range(queue_count)]
        
        # 优先级阈值，初始为默认值
        self.priority_thresholds = [1, 5, 10]
        
        # 优先级映射表 {task_id: priority}
        self.priority_map = {}
        
        # 最后一次重新平衡时间
        self.last_rebalance_time = datetime.now()
        
        # 优先级统计信息
        self.priority_stats = []
        
        logger.info(f"初始化优先级队列管理器: 队列数={queue_count}, 阈值百分比={self.threshold_percentiles}, "
                   f"工作线程比例={self.worker_ratios}")
    
    async def enqueue(self, task_id: str, priority: int) -> int:
        """
        将任务添加到适当的优先级队列
        
        Args:
            task_id: 任务ID
            priority: 任务优先级
            
        Returns:
            int: 队列索引
        """
        # 检查是否需要重新平衡
        await self._check_rebalance()
        
        # 确定任务应该进入哪个队列
        queue_index = self._get_queue_index(priority)
        
        # 保存优先级信息
        self.priority_map[task_id] = priority
        self.priority_stats.append(priority)
        
        # 添加到队列（使用负优先级值）
        await self.priority_queues[queue_index].put((-priority, task_id))
        logger.info(f"任务 {task_id} (优先级 {priority}) 添加到队列 {queue_index}")
        
        return queue_index
    
    async def dequeue(self, worker_id: int, worker_count: int) -> Tuple[int, Any]:
        """
        从适当的队列中获取任务
        
        Args:
            worker_id: 工作线程ID
            worker_count: 工作线程总数
            
        Returns:
            Tuple[int, Any]: (队列索引, 任务数据)
        """
        # 确定工作线程应该从哪个队列获取任务
        queue_index = self._get_worker_queue(worker_id, worker_count)
        # 检查指定队列是否有任务
        if not self.priority_queues[queue_index].empty():
            item = await self.priority_queues[queue_index].get()
            return queue_index, item
        
        # 如果指定队列为空，尝试从其他队列获取任务
        for i in range(self.queue_count):
            if not self.priority_queues[i].empty():
                item = await self.priority_queues[i].get()
                return i, item
        
        # 所有队列都为空，返回None
        return -1, None
    
    def task_done(self, queue_index: int) -> None:
        """
        标记任务完成
        
        Args:
            queue_index: 队列索引
        """
        if 0 <= queue_index < self.queue_count:
            self.priority_queues[queue_index].task_done()
    
    def remove_task(self, task_id: str) -> bool:
        """
        从优先级映射表中移除任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功移除
        """
        if task_id in self.priority_map:
            self.priority_map.pop(task_id)
            return True
        return False
    
    async def _check_rebalance(self) -> None:
        """
        检查是否需要重新平衡队列阈值
        """
        now = datetime.now()
        if (now - self.last_rebalance_time).total_seconds() > self.rebalance_interval:
            await self._rebalance_thresholds()
            self.last_rebalance_time = now
    
    async def _rebalance_thresholds(self) -> None:
        """
        重新平衡优先级阈值
        基于历史优先级分布动态调整阈值
        """
        if len(self.priority_stats) < 10:
            # 数据不足，使用默认值
            return
        
        # 复制并排序优先级数据
        sorted_priorities = sorted(self.priority_stats, reverse=True)
        
        # 根据百分比计算新阈值
        new_thresholds = []
        for percentile in self.threshold_percentiles[:-1]:  # 最后一个阈值不需要计算
            index = min(int(len(sorted_priorities) * percentile), len(sorted_priorities) - 1)
            new_thresholds.append(sorted_priorities[index])
        
        # 添加最小阈值（1，确保所有任务都能分配到队列）
        new_thresholds.append(1)
        
        # 更新阈值
        if new_thresholds != self.priority_thresholds:
            old_thresholds = self.priority_thresholds.copy()
            self.priority_thresholds = new_thresholds
            logger.info(f"优先级阈值更新: {old_thresholds} -> {new_thresholds}")
            
            # 限制历史数据大小，保留最近的数据
            if len(self.priority_stats) > 1000:
                self.priority_stats = self.priority_stats[-1000:]
    
    def _get_queue_index(self, priority: int) -> int:
        """
        根据优先级确定队列索引
        
        Args:
            priority: 任务优先级
            
        Returns:
            int: 队列索引
        """
        for i, threshold in enumerate(self.priority_thresholds):
            if priority >= threshold:
                return i
        
        # 默认放入最低优先级队列
        return self.queue_count - 1
    
    def _get_worker_queue(self, worker_id: int, worker_count: int) -> int:
        """
        根据工作线程ID确定应该从哪个队列获取任务
        
        Args:
            worker_id: 工作线程ID
            worker_count: 工作线程总数
            
        Returns:
            int: 队列索引
        """
        # 计算每个队列的工作线程数量
        worker_counts = [int(ratio * worker_count) for ratio in self.worker_ratios]
        
        # 补充剩余的工作线程到最后一个队列
        remaining = worker_count - sum(worker_counts)
        worker_counts[-1] += remaining
        
        # 确定工作线程应该处理哪个队列
        current_count = 0
        for i, count in enumerate(worker_counts):
            current_count += count
            if worker_id < current_count:
                return i
        
        # 默认返回最低优先级队列
        return self.queue_count - 1
    
    def get_queue_sizes(self) -> List[int]:
        """
        获取各队列大小
        
        Returns:
            List[int]: 队列大小列表
        """
        return [queue.qsize() for queue in self.priority_queues]
    
    def get_queue_status(self) -> Dict[str, Any]:
        """
        获取队列状态信息
        
        Returns:
            Dict[str, Any]: 队列状态
        """
        return {
            "queue_count": self.queue_count,
            "thresholds": self.priority_thresholds,
            "queue_sizes": self.get_queue_sizes(),
            "worker_ratios": self.worker_ratios,
            "task_count": len(self.priority_map),
            "last_rebalance": self.last_rebalance_time.strftime("%Y-%m-%d %H:%M:%S")
        } 