# agent/metrics_collector.py
import psutil
import time
from datetime import datetime

class NodeMetrics:
    """实时采集 VPS 节点的多维度指标"""
    
    def __init__(self, node_id: str):
        self.node_id = node_id
        self.window_seconds = 60  # 滑动窗口
        self.metrics_history = []
    
    def collect(self) -> dict:
        now = datetime.utcnow()
        
        # CPU 使用率（1分钟平均）
        cpu_percent = psutil.cpu_percent(interval=1)
        cpu_cores = psutil.cpu_count()
        
        # 内存使用
        mem = psutil.virtual_memory()
        
        # 磁盘 IO
        disk_io = psutil.disk_io_counters()
        
        # 网络 IO
        net_io = psutil.net_io_counters()
        
        # 系统负载
        load_avg = psutil.getloadavg()
        
        # 进程数
        processes = len(psutil.pids())
        
        return {
            "node_id": self.node_id,
            "timestamp": now.isoformat(),
            "cpu_percent": cpu_percent,
            "cpu_cores": cpu_cores,
            "memory_percent": mem.percent,
            "memory_available_mb": mem.available / 1024 / 1024,
            "disk_read_mb": disk_io.read_bytes / 1024 / 1024,
            "disk_write_mb": disk_io.write_bytes / 1024 / 1024,
            "net_sent_mb": net_io.bytes_sent / 1024 / 1024,
            "net_recv_mb": net_io.bytes_recv / 1024 / 1024,
            "load_1min": load_avg[0],
            "load_5min": load_avg[1],
            "load_15min": load_avg[2],
            "process_count": processes,
        }

# agent/health_score.py
class HealthScoreCalculator:
    """
    节点健康分 = f(cpu, memory, disk_io, latency, error_rate)
    范围: 0 (完全不可用) ~ 100 (最优状态)
    """
    
    WEIGHTS = {
        "cpu": 0.20,
        "memory": 0.20,
        "disk_io": 0.15,
        "latency": 0.25,
        "error_rate": 0.20,
    }
    
    def calculate(self, metrics: dict, baseline: dict) -> float:
        """基于当前指标和基线计算健康分"""
        scores = {}
        
        # CPU 健康分（越低越好）
        scores["cpu"] = max(0, 100 - metrics["cpu_percent"])
        
        # 内存健康分
        scores["memory"] = max(0, 100 - metrics["memory_percent"])
        
        # 磁盘 IO 压力分
        io_pressure = (metrics["disk_read_mb"] + metrics["disk_write_mb"]) / max(baseline.get("avg_disk_io", 1), 1)
        scores["disk_io"] = max(0, 100 - min(io_pressure * 30, 100))
        
        # 延迟健康分（由 API Gateway 上报）
        latency = metrics.get("avg_latency_ms", 50)
        scores["latency"] = max(0, 100 - (latency / 5))  # 250ms = 50分
        
        # 错误率健康分
        error_rate = metrics.get("error_rate_percent", 0)
        scores["error_rate"] = max(0, 100 - error_rate * 10)
        
        # 加权求和
        health_score = sum(scores[k] * self.WEIGHTS[k] for k in self.WEIGHTS)
        
        return round(health_score, 1)

# predictor/traffic_forecaster.py
from prophet import Prophet
import pandas as pd

class TrafficForecaster:
    """
    基于历史 QPS 数据预测未来 30 分钟流量
    """
    
    def __init__(self, window_hours=24):
        self.window_hours = window_hours
        self.models = {}  # 按路径/接口独立建模
    
    def fit(self, df: pd.DataFrame, path: str = "default"):
        """训练预测模型"""
        model = Prophet(
            yearly_seasonality=True,
            weekly_seasonality=True,
            daily_seasonality=True,
            changepoint_prior_scale=0.05,
        )
        model.fit(df)
        self.models[path] = model
        return model
    
    def predict(self, path: str = "default", hours_ahead=0.5) -> dict:
        """预测未来流量"""
        if path not in self.models:
            return {"forecast": [], "confidence": {}}
        
        model = self.models[path]
        future = model.make_future_dataframe(periods=int(hours_ahead * 24 * 60))
        forecast = model.predict(future)
        
        # 提取最近 N 个预测点
        recent = forecast.tail(int(hours_ahead * 12)).copy()
        
        return {
            "forecast": recent[["ds", "yhat", "yhat_lower", "yhat_upper"]].to_dict('records'),
            "predicted_peak_qps": float(recent["yhat"].max()),
            "predicted_avg_qps": float(recent["yhat"].mean()),
            "confidence_interval": {
                "lower": float(recent["yhat_lower"].min()),
                "upper": float(recent["yhat_upper"].max()),
            }
        }

# dispatcher/scheduler.py
import numpy as np
from scipy.optimize import linprog

class AIScheduler:
    """
    AI 智能调度器
    基于节点健康分 + 预测流量，动态计算最优路由权重
    """
    
    def __init__(self, nodes: list, weights: dict = None):
        self.nodes = nodes  # [{id, health_score, capacity, ...}, ...]
        self.weights = weights or {"latency": 0.4, "balance": 0.35, "cost": 0.25}
    
    def compute_weights(self, predicted_qps: float) -> dict:
        """
        计算每台 VPS 的最优权重
        """
        n = len(self.nodes)
        if n == 0:
            return {}
        
        # 构建优化变量：每台的权重 w_i (0 <= w_i <= 1)
        # 目标：最小化加权成本
        
        # 延迟成本：健康分低的节点延迟高
        latency_cost = np.array([
            (100 - node["health_score"]) / 100 * node["current_latency_ms"]
            for node in self.nodes
        ])
        
        # 负载均衡成本：偏离平均负载的程度
        avg_load = np.mean([node["cpu_percent"] for node in self.nodes])
        balance_cost = np.array([
            abs(node["cpu_percent"] - avg_load) / 100
            for node in self.nodes
        ])
        
        # 成本成本：按节点规格
        cost_per_node = np.array([
            node.get("hourly_cost", 0.05) for node in self.nodes
        ])
        
        # 组合目标函数
        c = (
            self.weights["latency"] * latency_cost +
            self.weights["balance"] * balance_cost +
            self.weights["cost"] * cost_per_node
        )
        
        # 约束：权重之和 = 1
        A_eq = np.ones((1, n))
        b_eq = np.array([1.0])
        
        # 边界：0 <= w_i <= 1
        bounds = [(0, 1)] * n
        
        # 求解
        result = linprog(c, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method='highs')
        
        if result.success:
            # 转换为百分比权重
            raw_weights = result.x
            total = raw_weights.sum()
            return {
                node["id"]: round(w / total * 100, 1)
                for node, w in zip(self.nodes, raw_weights)
                if w > 0.01  # 忽略极低权重的节点
            }
        else:
            # 回退到健康分加权
            return self._fallback_weighting()
    
    def _fallback_weighting(self) -> dict:
        """回退策略：按健康分比例分配"""
        total_health = sum(n["health_score"] for n in self.nodes)
        return {
            node["id"]: round(node["health_score"] / total_health * 100, 1)
            for node in self.nodes
        }

-- nginx/conf.d/ai_lb.lua
-- AI 调度器输出的权重实时更新到 Nginx upstream

local redis = require "resty.redis"
local cjson = require "cjson"

-- 连接 Redis（AI-LB 将权重写入 Redis）
local red = redis.new()
red:set_timeout(1000)
red:connect("127.0.0.1", 6379)

-- 获取最新权重配置
local weights_json = red:get("ai_lb:weights")
if weights_json then
    local weights = cjson.decode(weights_json)
    -- weights = {vps01=28, vps02=8, vps03=42, vps04=0}
    
    -- 动态设置 upstream 权重（通过 Lua 动态配置）
    -- 实际生产环境建议使用 Nginx Plus API 或 Envoy xDS
end

# dispatcher/envoy_xds_client.py
from envoy_control import XdsClient  # 简化示意

class EnvoyDispatcher:
    """通过 xDS 协议动态更新 Envoy 路由配置"""
    
    def __init__(self, xds_host="127.0.0.1", xds_port=5678):
        self.client = XdsClient(host=xds_host, port=xds_port)
    
    def update_endpoints(self, weights: dict):
        """
        更新后端端点权重
        weights = {"vps-01": 28, "vps-02": 8, "vps-03": 42}
        """
        endpoints = []
        for node_id, weight in weights.items():
            endpoints.append({
                "address": f"{node_id}.internal",
                "port": 8080,
                "weight": weight,
                "health_check": weight > 0,  # 权重为 0 时标记不健康
            })
        
        self.client.update_cluster("backend_cluster", endpoints)
        print(f"Updated {len(endpoints)} endpoints via xDS")

# verifier/result_evaluator.py
class ResultEvaluator:
    """
    验证调度决策的实际效果
    使用 A/B 对比：调度前 5min 平均 vs 调度后 5min 平均
    """
    
    def evaluate(self, before: dict, after: dict) -> dict:
        metrics = {
            "avg_latency_ms": {
                "before": before["avg_latency_ms"],
                "after": after["avg_latency_ms"],
                "change_pct": self._pct_change(before["avg_latency_ms"], after["avg_latency_ms"]),
            },
            "error_rate": {
                "before": before["error_rate_percent"],
                "after": after["error_rate_percent"],
                "change_pct": self._pct_change(before["error_rate_percent"], after["error_rate_percent"]),
            },
            "load_variance": {
                "before": before["load_variance"],
                "after": after["load_variance"],
                "change_pct": self._pct_change(before["load_variance"], after["load_variance"]),
            },
            "decision_valid": after["avg_latency_ms"] < before["avg_latency_ms"],
        }
        return metrics
    
    def _pct_change(self, old, new):
        if old == 0:
            return 0
        return round((new - old) / old * 100, 1)

ai-load-balancer/
├── agent/                  # VPS 节点采集器
│   ├── metrics_collector.py
│   ├── health_score.py
│   └── agent.sh           # 启动脚本
├── predictor/              # 流量预测
│   ├── forecaster.py
│   └── models/            # 保存的训练模型
├── dispatcher/             # 调度决策
│   ├── scheduler.py
│   ├── envoy_xds_client.py
│   └── nginx_lua/
├── verifier/               # 效果验证
│   └── result_evaluator.py
├── config/
│   ├── nodes.yaml         # VPS 清单
│   ├── weights.yaml       # 优化权重配置
│   └── thresholds.yaml    # 告警阈值
├── orchestrator.py         # 主入口（定时调度）
└── requirements.txt

# orchestrator.py
#!/usr/bin/env python3
"""
AI Load Balancer Orchestrator
每 30 秒执行一次完整调度循环
"""

import time
import yaml
from datetime import datetime
from agent.metrics_collector import NodeMetrics
from agent.health_score import HealthScoreCalculator
from predictor.traffic_forecaster import TrafficForecaster
from dispatcher.scheduler import AIScheduler
from dispatcher.envoy_xds_client import EnvoyDispatcher
from verifier.result_evaluator import ResultEvaluator

def main():
    # 加载配置
    with open("config/nodes.yaml") as f:
        config = yaml.safe_load(f)
    
    # 初始化组件
    calculators = {node["id"]: HealthScoreCalculator() for node in config["nodes"]}
    forecaster = TrafficForecaster()
    scheduler = AIScheduler(config["nodes"])
    dispatcher = EnvoyDispatcher()
    evaluator = ResultEvaluator()
    
    print(f"[{datetime.utcnow().isoformat()}] AI-LB started, monitoring {len(config['nodes'])} nodes")
    
    while True:
        try:
            # Step 1: 采集所有节点指标
            all_metrics = {}
            for node in config["nodes"]:
                nid = node["id"]
                collector = NodeMetrics(nid)
                all_metrics[nid] = collector.collect()
            
            # Step 2: 计算健康分
            health_scores = {}
            for nid, metrics in all_metrics.items():
                health_scores[nid] = calculators[nid].calculate(metrics, {})
            
            # Step 3: 预测未来流量
            prediction = forecaster.predict(hours_ahead=0.5)
            
            # Step 4: 计算最优权重
            weights = scheduler.compute_weights(prediction.get("predicted_peak_qps", 0))
            
            # Step 5: 下发配置
            dispatcher.update_endpoints(weights)
            
            # Step 6: 记录效果
            print(f"[{datetime.utcnow().isoformat()}] Weights: {weights} | "
                  f"Predicted QPS: {prediction.get('predicted_peak_qps', 'N/A')}")
            
            # 等待下一个周期
            time.sleep(30)
            
        except Exception as e:
            print(f"[ERROR] {e}")
            time.sleep(60)

if __name__ == "__main__":
    main()

# 1. 安装依赖
pip install prophet scipy psutil pyyaml redis envoy-control

# 2. 在每个 VPS 上部署 Agent
scp agent/agent.sh user@vps-01:~/ai-lb-agent/
ssh user@vps-01 '~/ai-lb-agent/agent.sh start'

# 3. 在主控制节点部署 Orchestrator
cd ~/ai-load-balancer
python3 orchestrator.py &

# 4. 配置 Envoy 接入 xDS（或使用 Nginx Lua 方案）
# Envoy 配置文件中添加 xDS 来源：
# discovery_service:
#   ads:
#     address: localhost:5678