class EnvelopeSLOTracker:
    """
    统一跟踪TTFT和TPOT的SLO达成情况
    使用松弛时间作为统一的度量标准
    """
    def __init__(self, ttft_slo_ms: float, tpot_slo_ms: float, gpu_profile: GPUProfile):
        self.ttft_slo = ttft_slo_ms
        self.tpot_slo = tpot_slo_ms
        self.gpu_profile = gpu_profile
        self.request_progress = {}  # request_id -> progress_info
        self.token_processing_times = deque(maxlen=1000)  # 用于TPOT预测的历史数据
        
    def calculate_slack(self, request_id: str, current_time: float) -> float:
        """
        计算请求的slack时间
        正值表示还有余量，负值表示已违规
        """
        if request_id not in self.request_progress:
            return float('inf')  # 新请求，假设无限slack
            
        req_info = self.request_progress[request_id]
        
        if req_info.is_prefill:
            # 预填充任务的slack计算
            elapsed = current_time - req_info.arrival_time
            remaining_tokens = req_info.total_tokens - req_info.processed_tokens
            
            # 使用GPU性能模型估计剩余处理时间
            estimated_prefill_time = self.estimate_prefill_time(
                remaining_tokens, 
                req_info.context_length
            )
            
            # 对于预填充任务，还需要考虑后续解码时间
            # 因为TTFT SLO包含整个预填充完成时间
            estimated_completion = elapsed + estimated_prefill_time
            slack = self.ttft_slo - estimated_completion
            
        else:
            # 解码任务的slack计算
            if req_info.processed_tokens == 0:
                # 还未开始处理，计算TTFT slack
                elapsed = current_time - req_info.arrival_time
                estimated_first_token_time = self.estimate_first_token_time(req_info)
                slack = self.ttft_slo - (elapsed + estimated_first_token_time)
            elif req_info.processed_tokens == 1:
                # 第一个token刚完成，仍然关注TTFT
                elapsed = current_time - req_info.arrival_time
                slack = self.ttft_slo - elapsed
            else:
                # 多个token已完成，关注TPOT
                # 使用历史数据预测下一个token的处理时间
                predicted_tpot = self.predict_next_tpot(req_info)
                expected_next_completion = current_time + predicted_tpot
                actual_tpot = expected_next_completion - req_info.last_token_time
                slack = self.tpot_slo - actual_tpot
        
        return slack
    
    def estimate_prefill_time(self, tokens: int, context_length: int) -> float:
        """基于GPU性能模型估计预填充时间"""
        # 简化的性能模型：考虑tokens数量和上下文长度
        base_time_per_token = self.gpu_profile.prefill_time_per_token
        context_factor = 1.0 + (context_length / 1000.0) * 0.1  # 上下文长度影响
        return tokens * base_time_per_token * context_factor
    
    def predict_next_tpot(self, req_info) -> float:
        """基于历史数据预测下一个TPOT"""
        if len(req_info.token_times) < 2:
            return self.gpu_profile.avg_decode_time
        
        # 使用最近几个token的平均时间作为预测
        recent_tpot = np.mean(req_info.token_times[-5:])
        return recent_tpot
    
    def update_progress(self, request_id: str, tokens_generated: int, processing_time: float):
        """更新请求进度和性能数据"""
        if request_id not in self.request_progress:
            return
            
        req_info = self.request_progress[request_id]
        req_info.processed_tokens += tokens_generated
        req_info.last_token_time = time.time()
        
        # 记录处理时间用于后续预测
        if tokens_generated == 1:
            req_info.token_times.append(processing_time)
            self.token_processing_times.append(processing_time)
        
        # 更新上下文长度
        req_info.context_length += tokens_generated

def compute_adaptive_batch_capacity(
    active_requests: List[Request],
    pending_requests: List[Request],
    slo_tracker: EnvelopeSLOTracker,
    gpu_profile: GPUProfile,
    safety_margin: float = 0.1
) -> Tuple[int, Dict[str, Any]]:
    """
    基于slack分析动态计算批容量
    返回容量和详细的分析报告
    """
    current_time = time.time()
    
    # 分别计算活跃请求和等待请求的slack
    active_slacks = []
    pending_slacks = []
    
    for req in active_requests:
        slack = slo_tracker.calculate_slack(req.id, current_time)
        active_slacks.append(slack)
        
    for req in pending_requests:
        slack = slo_tracker.calculate_slack(req.id, current_time)
        pending_slacks.append(slack)
    
    # 计算关键统计指标
    if active_slacks:
        active_avg_slack = np.mean(active_slacks)
        active_p95_slack = np.percentile(active_slacks, 95)
        active_p5_slack = np.percentile(active_slacks, 5)
    else:
        active_avg_slack = float('inf')
        active_p95_slack = float('inf')
        active_p5_slack = float('inf')
        
    if pending_slacks:
        pending_avg_slack = np.mean(pending_slacks)
        pending_min_slack = min(pending_slacks)
    else:
        pending_avg_slack = float('inf')
        pending_min_slack = float('inf')
    
    # 基础容量设置
    base_capacity = gpu_profile.max_batch_size
    
    # 决策逻辑：综合考虑多个因素
    decision_factors = {
        'active_avg_slack': active_avg_slack,
        'active_p5_slack': active_p5_slack,  # 关注最紧张的请求
        'pending_min_slack': pending_min_slack,  # 关注最紧急的等待请求
        'pending_count': len(pending_requests),
        'active_count': len(active_requests)
    }
    
    # 主要决策逻辑
    if pending_min_slack < -50:  # 有请求即将严重违规
        # 紧急模式：最小化批容量以快速响应
        capacity = max(1, int(base_capacity * 0.3))
        mode = "emergency"
    elif active_p5_slack < -20:  # 有活跃请求时间紧张
        # 保守模式：减少容量以保护SLO
        capacity = int(base_capacity * 0.6)
        mode = "conservative"
    elif pending_avg_slack > 100 and len(pending_requests) > 10:  # 有大量宽松的等待请求
        # 扩展模式：增加容量以提高吞吐量
        capacity_increase = min(int(active_avg_slack / 20), 8)
        capacity = min(base_capacity + capacity_increase, gpu_profile.hard_limit)
        mode = "expansion"
    elif active_avg_slack > 50:  # 系统整体宽松
        # 宽松模式：适度增加容量
        capacity = min(base_capacity + 3, gpu_profile.hard_limit)
        mode = "relaxed"
    else:
        # 正常模式：使用基准容量
        capacity = base_capacity
        mode = "normal"
    
    # 应用安全边际
    capacity = int(capacity * (1 - safety_margin))
    
    # 确保最小容量
    capacity = max(capacity, 1)
    
    analysis_report = {
        'mode': mode,
        'factors': decision_factors,
        'final_capacity': capacity
    }
    
    return capacity, analysis_report

def get_urgent_decode_requests(self, decode_requests: List[Request], 
                             current_time: float, violation_threshold: float) -> List[Request]:
    """获取濒临SLO违规的解码请求"""
    urgent_requests = []
    for req in decode_requests:
        slack = self.slo_tracker.calculate_slack(req.id, current_time)
        if slack <= violation_threshold:
            # 计算违规的严重程度
            urgency_score = -slack  # 负slack越大，越紧急
            urgent_requests.append((req, urgency_score, slack))
    
    # 按紧急程度排序，最紧急的优先
    urgent_requests.sort(key=lambda x: x[1], reverse=True)
    return [(req, slack) for req, _, slack in urgent_requests]

def select_prefill_chunk(self, pending_prefill: List[Request], 
                       available_capacity: int, current_time: float) -> List[Request]:
    """选择合适的预填充任务块"""
    if not pending_prefill or available_capacity <= 0:
        return []
    
    # 计算每个预填充请求的紧急程度
    prefill_candidates = []
    for req in pending_prefill:
        slack = self.slo_tracker.calculate_slack(req.id, current_time)
        urgency = -slack if slack < 0 else 0  # 只有违规的才紧急
        cost = self.estimate_prefill_cost(req)
        
        # 如果容量不足以处理整个请求，考虑chunking
        if cost > available_capacity:
            # 创建chunk
            chunk_size = min(available_capacity, req.remaining_tokens)
            chunk_cost = self.estimate_chunk_cost(chunk_size, req.context_length)
            if chunk_cost <= available_capacity:
                chunk = self.create_prefill_chunk(req, chunk_size)
                prefill_candidates.append((chunk, urgency, chunk_cost))
        else:
            prefill_candidates.append((req, urgency, cost))
    
    # 按紧急程度排序
    prefill_candidates.sort(key=lambda x: x[1], reverse=True)
    
    # 贪心选择
    selected = []
    current_load = 0
    for candidate, urgency, cost in prefill_candidates:
        if current_load + cost <= available_capacity:
            selected.append(candidate)
            current_load += cost
        
    return selected

def select_non_urgent_decode(self, active_decode: List[Request], 
                           remaining_capacity: int, current_time: float) -> List[Request]:
    """选择非紧急的解码任务填充剩余容量"""
    non_urgent = []
    for req in active_decode:
        slack = self.slo_tracker.calculate_slack(req.id, current_time)
        if slack > -10:  # 不是紧急任务
            cost = self.estimate_decode_cost(req)
            non_urgent.append((req, slack, cost))
    
    # 按slack排序，slack大的（更宽松的）优先
    non_urgent.sort(key=lambda x: x[1], reverse=True)
    
    selected = []
    current_load = 0
    for req, slack, cost in non_urgent:
        if current_load + cost <= remaining_capacity:
            selected.append(req)
            current_load += cost
    
    return selected

def form_batch(self, pending_prefill: List[Request], 
               active_decode: List[Request],
               current_capacity: int) -> Batch:
    batch = Batch()
    current_time = time.time()
    
    # Phase 1: 优先处理濒临SLO违规的解码请求
    urgent_decode = self.get_urgent_decode_requests(
        active_decode, current_time, violation_threshold=-10
    )
    
    # 添加紧急解码任务
    urgent_capacity_used = 0
    for req, slack in urgent_decode:
        cost = self.estimate_decode_cost(req)
        if urgent_capacity_used + cost <= current_capacity:
            batch.add_request(req)
            urgent_capacity_used += cost
        else:
            break
    
    # Phase 2: 立即调度预填充请求
    available_capacity = current_capacity - urgent_capacity_used
    if available_capacity > 0:
        prefill_chunk = self.select_prefill_chunk(
            pending_prefill, available_capacity, current_time
        )
        for chunk in prefill_chunk:
            batch.add_request(chunk)
    
    # Phase 3: 填充剩余容量
    remaining = current_capacity - batch.get_total_cost()
    if remaining > 0:
        non_urgent = self.select_non_urgent_decode(
            active_decode, remaining, current_time
        )
        for req in non_urgent:
            batch.add_request(req)
    
    return batch

# 传统调度器的优化目标
# maximize throughput = f(batch_size, gpu_utilization)

class TraditionalScheduler:
    def __init__(self, max_batch_size):
        self.max_batch_size = max_batch_size
        
    def schedule(self, requests):
        # 简单的贪心策略：尽可能填满批次
        batch = []
        current_size = 0
        for req in sorted(requests, key=lambda x: x.arrival_time):
            if current_size + req.cost <= self.max_batch_size:
                batch.append(req)
                current_size += req.cost
        return batch

# FairBatching的多目标优化
# minimize max_violation + λ * (1 - gpu_utilization)
# subject to: all_slo_constraints

class FairBatchingScheduler:
    def __init__(self, slo_tracker, gpu_profile):
        self.slo_tracker = slo_tracker
        self.gpu_profile = gpu_profile
        self.lambda_param = 0.1  # 权衡参数
        
    def schedule(self, requests):
        # 复杂的多阶段决策过程
        pending_prefill = [r for r in requests if r.is_prefill and not r.started]
        active_decode = [r for r in requests if not r.is_prefill and r.started]
        
        # 动态分析当前状态
        current_capacity, analysis = compute_adaptive_batch_capacity(
            active_decode, pending_prefill, self.slo_tracker, self.gpu_profile
        )
        
        # 基于SLO跟踪的三阶段决策
        batch = self.form_batch(pending_prefill, active_decode, current_capacity)
        
        return batch

# vLLM scheduler.py 中的集成示例
from vllm.core.scheduler import Scheduler, SchedulingPolicy
from vllm.core.interfaces import RequestMetadata
import time

class FairBatchingRequest(RequestMetadata):
    """扩展的请求元数据，支持FairBatching所需的状态跟踪"""
    def __init__(self, original_request, ttft_slo_ms, tpot_slo_ms):
        super().__init__(original_request)
        self.arrival_time = time.time()
        self.ttft_slo_ms = ttft_slo_ms
        self.tpot_slo_ms = tpot_slo_ms
        self.processed_tokens = 0
        self.last_token_time = self.arrival_time
        self.token_times = []
        self.context_length = len(original_request.prompt_token_ids)
        self.is_prefill = True
        self.slack_cache = None
        self.slack_cache_time = 0
        
    def mark_as_decode(self):
        """标记请求进入解码阶段"""
        self.is_prefill = False
        
    def update_token_stats(self, token_time):
        """更新token处理统计信息"""
        self.processed_tokens += 1
        self.token_times.append(token_time)
        self.last_token_time = time.time()

class FairBatchingScheduler(Scheduler):
    def __init__(self, scheduler_config, cache_config, lora_config=None):
        super().__init__(scheduler_config, cache_config, lora_config)
        
        # FairBatching特定配置
        self.ttft_slo_ms = getattr(scheduler_config, 'ttft_slo_ms', 500.0)
        self.tpot_slo_ms = getattr(scheduler_config, 'tpot_slo_ms', 100.0)
        self.violation_threshold_ms = getattr(scheduler_config, 'violation_threshold_ms', -10.0)
        self.adaptive_scaling_enabled = getattr(scheduler_config, 'adaptive_scaling_enabled', True)
        
        # 初始化SLO跟踪器
        self.slo_tracker = EnvelopeSLOTracker(
            ttft_slo_ms=self.ttft_slo_ms,
            tpot_slo_ms=self.tpot_slo_ms,
            gpu_profile=self._build_gpu_profile()
        )
        
        # FairBatching策略
        self.policy = FairBatchingPolicy(
            slo_tracker=self.slo_tracker,
            violation_threshold=self.violation_threshold_ms
        )
        
        # 请求映射表
        self.request_map = {}
    
    def add_seq_group(self, seq_group):
        """添加新的序列组，创建FairBatching请求对象"""
        fb_request = FairBatchingRequest(
            seq_group, 
            self.ttft_slo_ms, 
            self.tpot_slo_ms
        )
        self.request_map[seq_group.request_id] = fb_request
        super().add_seq_group(seq_group)
    
    def schedule(self):
        """FairBatching调度主逻辑"""
        # 获取等待队列和运行队列
        waiting_queue = []
        running_queue = []
        
        # 转换vLLM内部表示为FairBatching请求
        for seq_group in self.waiting:
            if seq_group.request_id in self.request_map:
                waiting_queue.append(self.request_map[seq_group.request_id])
                
        for seq_group in self.running:
            if seq_group.request_id in self.request_map:
                req = self.request_map[seq_group.request_id]
                # 检查是否已完成预填充
                if seq_group.is_prefill_complete():
                    req.mark_as_decode()
                running_queue.append(req)
        
        # 计算自适应批容量
        if self.adaptive_scaling_enabled:
            current_capacity, _ = compute_adaptive_batch_capacity(
                running_queue, waiting_queue, self.slo_tracker, 
                self.slo_tracker.gpu_profile
            )
        else:
            current_capacity = self.scheduler_config.max_num_seqs
        
        # 应用FairBatching策略
        scheduled_requests = self.policy.form_batch(
            waiting_queue, running_queue, current_capacity
        )
        
        # 转换回vLLM格式
        scheduled_seq_groups = []
        for req in scheduled_requests:
            # 找到对应的seq_group
            for seq_group in self.waiting + self.running:
                if seq_group.request_id == req.request_id:
                    scheduled_seq_groups.append(seq_group)
                    break
        
        # 调用父类的调度逻辑
        return super().schedule_from_selected(scheduled_seq_groups)
    
    def _build_gpu_profile(self):
        """基于当前GPU配置构建性能模型"""
        # 这里需要根据实际的GPU型号和配置进行调整
        # 可以通过基准测试自动校准
        return GPUProfile(
            max_batch_size=self.scheduler_config.max_num_seqs,
            hard_limit=self.scheduler_config.max_num_seqs * 1.5,
            prefill_time_per_token=0.5,  # ms per token
            avg_decode_time=20.0,  # ms per token
            memory_bandwidth_gb=1555  # GB/s for A100
        )

class FairBatchingPolicy(SchedulingPolicy):
    def __init__(self, slo_tracker, violation_threshold=-10.0):
        self.slo_tracker = slo_tracker
        self.violation_threshold = violation_threshold
    
    def form_batch(self, waiting_queue, running_queue, capacity):
        # 实现三阶段批打包逻辑
        # ... (前面详细实现的代码)
        pass

# vLLM配置扩展
@dataclass
class SchedulerConfig:
    # 现有配置...
    max_num_seqs: int = 256
    max_model_len: int = 2048
    
    # FairBatching扩展配置（可选）
    enable_fair_batching: bool = False
    ttft_slo_ms: Optional[float] = None
    tpot_slo_ms: Optional[float] = None
    violation_threshold_ms: float = -10.0
    adaptive_scaling_enabled: bool = True
    
    def __post_init__(self):
        if self.enable_fair_batching:
            if self.ttft_slo_ms is None:
                self.ttft_slo_ms = 500.0  # 默认500ms
            if self.tpot_slo_ms is None:
                self.tpot_slo_ms = 100.0  # 默认100ms

class OnlineTuner:
    def __init__(self, scheduler):
        self.scheduler = scheduler
        self.slo_violation_history = deque(maxlen=1000)
        
    def adjust_parameters(self, current_metrics):
        """基于当前指标动态调整参数"""
        violation_rate = current_metrics['slo_violation_rate']
        
        if violation_rate > 0.05:  # 违规率过高
            # 收紧SLO或降低批容量
            self.scheduler.ttft_slo_ms *= 0.9
            self.scheduler.adaptive_scaling_enabled = True
        elif violation_rate < 0.01:  # 违规率过低
            # 放宽SLO以提高吞吐量
            self.scheduler.ttft_slo_ms *= 1.1

# 监控指标示例
metrics = {
    'ttft_p99_ms': 520,
    'ttft_slo_achievement_rate': 0.98,
    'tpot_p99_ms': 42,
    'tpot_slo_achievement_rate': 0.99,
    'gpu_utilization': 0.82,
    'fairness_cv': 0.23,
    'current_batch_capacity': 180,
    'scheduling_mode': 'conservative'
}

class MLEnhancedFairBatching(FairBatchingScheduler):
    def __init__(self, model_path: str, feature_config: Dict):
        super().__init__()
        self.feature_extractor = FeatureExtractor(feature_config)
        self.behavior_predictor = load_ml_model(model_path)
        self.prediction_cache = LRUCache(maxsize=10000)
    
    def predict_request_behavior(self, request: Request) -> Dict[str, float]:
        """预测请求的行为特征"""
        cache_key = request.request_id
        if cache_key in self.prediction_cache:
            return self.prediction_cache[cache_key]
            
        features = self.feature_extractor.extract(request)
        predictions = self.behavior_predictor.predict(features)
        
        # 缓存预测结果
        self.prediction_cache[cache_key] = predictions
        return predictions
    
    def calculate_slack(self, request_id: str, current_time: float) -> float:
        """使用ML预测增强的slack计算"""
        predictions = self.predict_request_behavior(
            self.request_map[request_id]
        )
        
        # 使用预测结果改进时间估计
        if predictions['is_long_running']:
            # 长运行请求可能需要特殊处理
            return self._calculate_slack_with_prediction(
                request_id, current_time, predictions
            )
        else:
            return super().calculate_slack(request_id, current_time)

class MultiTenantFairBatching(FairBatchingScheduler):
    def __init__(self, tenant_configs: Dict[str, TenantConfig]):
        super().__init__()
        self.tenant_configs = tenant_configs
        self.tenant_slo_trackers = {
            tenant_id: EnvelopeSLOTracker(
                config.ttft_slo_ms, config.tpot_slo_ms
            )
            for tenant_id, config in tenant_configs.items()
        }
        
    def form_batch(self, waiting_queue, running_queue, capacity):
        # 按租户分组请求
        tenant_requests = defaultdict(lambda: {'waiting': [], 'running': []})
        for req in waiting_queue:
            tenant_requests[req.tenant_id]['waiting'].append(req)
        for req in running_queue:
            tenant_requests[req.tenant_id]['running'].append(req)
            
        # 为每个租户分配容量
        tenant_capacities = self._allocate_capacity_by_tenant(
            tenant_requests, capacity
        )
        
        # 分别调度每个租户的请求
        final_batch = []
        for tenant_id, capacities in tenant_capacities.items():
            tenant_batch = self._schedule_tenant_requests(
                tenant_requests[tenant_id]['waiting'],
                tenant_requests[tenant_id]['running'],
                capacities
            )
            final_batch.extend(tenant_batch)
            
        return final_batch