# 安装 Ollama
curl -fsSL https://ollama.com/install.sh | sh

# 拉取轻量级模型（适合 VPS 资源）
ollama pull llama3.2:3b
ollama pull nomic-embed-text  # 用于 RAG 检索

# 验证安装
ollama list

pip install langgraph langchain langchain-community \
            langchain-ollama prometheus-client \
            python-dotenv docker requests

# 创建项目结构
mkdir -p ~/vps-selfheal/{agents,tools,knowledge}
cd ~/vps-selfheal

# config.yaml
llm:
  model: "llama3.2:3b"
  base_url: "http://localhost:11434"
  
monitoring:
  prometheus_url: "http://localhost:9090"
  check_interval: 60  # 秒
  
alert_thresholds:
  cpu_percent: 85
  memory_percent: 80
  disk_percent: 90
  container_restart_count: 3

notification:
  provider: "telegram"
  bot_token: "${TELEGRAM_BOT_TOKEN}"
  chat_id: "${TELEGRAM_CHAT_ID}"

# agents/monitor.py
import requests
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator

class MonitorState(TypedDict):
    timestamp: str
    metrics: dict
    anomalies: list
    severity: str  # "normal", "warning", "critical"

class MonitorAgent:
    def __init__(self, config: dict):
        self.config = config
        self.prometheus_url = config["monitoring"]["prometheus_url"]
        
    def query_prometheus(self, query: str) -> float:
        """查询 Prometheus 指标"""
        url = f"{self.prometheus_url}/api/v1/query"
        params = {"query": query}
        resp = requests.get(url, params=params, timeout=10)
        data = resp.json()
        if data["status"] == "success" and data["data"]["result"]:
            return float(data["data"]["result"][0]["value"][1])
        return 0.0
    
    def collect_metrics(self) -> dict:
        """采集系统关键指标"""
        return {
            "cpu_percent": self.query_prometheus(
                "100 - (avg(irate(node_cpu_seconds_total{mode=\"idle\"}[5m])) * 100)"
            ),
            "memory_percent": self.query_prometheus(
                "(1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes)) * 100"
            ),
            "disk_percent": self.query_prometheus(
                "(1 - node_filesystem_avail_bytes{mountpoint=\"/\"} / node_filesystem_size_bytes{mountpoint=\"/\"}) * 100"
            ),
            "container_restarts": self.query_prometheus(
                "sum(increase(container_restart_count[10m]))"
            ),
            "load_avg_1m": self.query_prometheus(
                "node_load1"
            ),
        }
    
    def detect_anomalies(self, metrics: dict) -> list:
        """检测异常"""
        thresholds = self.config["alert_thresholds"]
        anomalies = []
        
        if metrics["cpu_percent"] > thresholds["cpu_percent"]:
            anomalies.append({
                "type": "high_cpu",
                "value": metrics["cpu_percent"],
                "threshold": thresholds["cpu_percent"],
                "message": f"CPU 使用率过高: {metrics['cpu_percent']:.1f}%"
            })
        
        if metrics["memory_percent"] > thresholds["memory_percent"]:
            anomalies.append({
                "type": "high_memory",
                "value": metrics["memory_percent"],
                "threshold": thresholds["memory_percent"],
                "message": f"内存使用率过高: {metrics['memory_percent']:.1f}%"
            })
        
        if metrics["disk_percent"] > thresholds["disk_percent"]:
            anomalies.append({
                "type": "high_disk",
                "value": metrics["disk_percent"],
                "threshold": thresholds["disk_percent"],
                "message": f"磁盘使用率过高: {metrics['disk_percent']:.1f}%"
            })
        
        if metrics["container_restarts"] > thresholds["container_restart_count"]:
            anomalies.append({
                "type": "container_crash",
                "value": metrics["container_restart_count"],
                "threshold": thresholds["container_restart_count"],
                "message": f"容器频繁重启: {metrics['container_restart_count']:.0f} 次/10分钟"
            })
        
        return anomalies
    
    def determine_severity(self, anomalies: list) -> str:
        """判断告警级别"""
        if not anomalies:
            return "normal"
        types = {a["type"] for a in anomalies}
        if "container_crash" in types or "high_memory" in types:
            return "critical"
        return "warning"
    
    def execute(self, state: MonitorState) -> MonitorState:
        """Agent 执行逻辑"""
        metrics = self.collect_metrics()
        anomalies = self.detect_anomalies(metrics)
        severity = self.determine_severity(anomalies)
        
        return {
            **state,
            "metrics": metrics,
            "anomalies": anomalies,
            "severity": severity,
        }

# agents/diagnose.py
from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
import json

class DiagnoseAgent:
    def __init__(self, llm: OllamaLLM, knowledge_base: list):
        self.llm = llm
        self.kb = knowledge_base  # 历史故障知识库
        
    def build_prompt(self, anomalies: list, metrics: dict) -> str:
        """构建诊断提示词"""
        kb_context = "\n".join([
            f"- {item['incident']}: {item['root_cause']} → {item['solution'}"
            for item in self.kb[-5:]  # 取最近5条
        ])
        
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是 VPS 运维专家。根据监控异常和系统指标，分析根因。
历史故障案例参考：
{kb_context}

请输出 JSON 格式：
{{
  "root_cause": "根因描述",
  "confidence": 0.0-1.0,
  "suggested_fixes": ["修复方案1", "修复方案2"],
  "risk_level": "low|medium|high",
  "requires_manual_review": true|false
}}"""),
            ("user", """当前异常：{anomalies}
当前指标：{metrics}"""),
        ])
        
        return prompt.format(
            kb_context=kb_context,
            anomalies=json.dumps(anomalies, ensure_ascii=False),
            metrics=json.dumps({k: round(v, 2) for k, v in metrics.items()}),
        )
    
    def execute(self, state: dict) -> dict:
        """执行诊断"""
        prompt = self.build_prompt(state["anomalies"], state["metrics"])
        response = self.llm.invoke(prompt)
        
        # 解析 LLM 输出
        try:
            diagnosis = json.loads(response.content)
        except json.JSONDecodeError:
            diagnosis = {
                "root_cause": "无法自动诊断，需要人工介入",
                "confidence": 0.0,
                "suggested_fixes": ["联系运维人员"],
                "risk_level": "high",
                "requires_manual_review": True,
            }
        
        state["diagnosis"] = diagnosis
        return state

# agents/fix.py
import subprocess
import docker
from datetime import datetime

class FixAgent:
    def __init__(self, config: dict):
        self.config = config
        self.docker_client = docker.from_env()
        
    def fix_high_cpu(self) -> dict:
        """处理高 CPU 异常"""
        results = []
        
        # 找出 CPU 占用最高的进程
        try:
            result = subprocess.run(
                ["ps", "-eo", "pid,pcpu,comm", "--sort=-pcpu", "--no-headers"],
                capture_output=True, text=True, timeout=10
            )
            top_process = result.stdout.split("\n")[0]
            pid = top_process.split()[0]
            cmd = " ".join(top_process.split()[2:])
            
            results.append({
                "action": "identify_top_process",
                "pid": pid,
                "command": cmd,
                "cpu": top_process.split()[1],
            })
            
            # 如果是异常进程，考虑终止
            if "java" in cmd or "python" in cmd:
                # 不自动 kill，记录告警
                results.append({
                    "action": "log_alert",
                    "message": f"高 CPU 进程: PID={pid}, CMD={cmd}",
                })
                
        except Exception as e:
            results.append({"action": "error", "message": str(e)})
            
        return {"success": True, "actions": results}
    
    def fix_high_memory(self) -> dict:
        """处理高内存异常"""
        results = []
        
        # 清理缓存
        try:
            subprocess.run(["sync"], timeout=5)
            subprocess.run(["sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], timeout=5)
            results.append({"action": "clear_page_cache", "success": True})
        except Exception as e:
            results.append({"action": "clear_page_cache", "success": False, "error": str(e)})
        
        # 找出内存占用最高的容器
        try:
            containers = self.docker_client.containers.list()
            mem_stats = []
            for c in containers:
                stats = c.stats(stream=False)
                mem_usage = stats.get("memory_stats", {}).get("usage", 0)
                mem_limit = stats.get("memory_stats", {}).get("limit", 1)
                mem_stats.append({
                    "name": c.name,
                    "id": c.short_id,
                    "memory_mb": mem_usage / 1024 / 1024,
                    "memory_limit_mb": mem_limit / 1024 / 1024,
                })
            mem_stats.sort(key=lambda x: x["memory_mb"], reverse=True)
            results.append({"action": "top_containers_by_memory", "data": mem_stats[:5]})
        except Exception as e:
            results.append({"action": "list_containers", "success": False, "error": str(e)})
            
        return {"success": True, "actions": results}
    
    def fix_container_crash(self) -> dict:
        """处理容器频繁重启"""
        results = []
        
        try:
            # 获取重启次数最多的容器
            containers = self.docker_client.containers.list(all=True)
            crash_info = []
            for c in containers:
                if c.attrs.get("State", {}).get("RestartCount", 0) > 0:
                    crash_info.append({
                        "name": c.name,
                        "restart_count": c.attrs["State"]["RestartCount"],
                        "status": c.attrs["State"]["Status"],
                        "image": c.image.tags[0] if c.image.tags else c.image.id[:12],
                    })
            
            results.append({"action": "identify_crashing_containers", "data": crash_info})
            
            # 自动重启故障容器
            for info in crash_info:
                try:
                    container = self.docker_client.containers.get(info["name"])
                    container.restart()
                    results.append({
                        "action": "restart_container",
                        "container": info["name"],
                        "success": True,
                    })
                except Exception as e:
                    results.append({
                        "action": "restart_container",
                        "container": info["name"],
                        "success": False,
                        "error": str(e),
                    })
        except Exception as e:
            results.append({"action": "error", "message": str(e)})
            
        return {"success": True, "actions": results}
    
    def fix_high_disk(self) -> dict:
        """处理磁盘空间不足"""
        results = []
        
        # 清理 Docker 垃圾
        try:
            result = subprocess.run(
                ["docker", "system", "prune", "-f"],
                capture_output=True, text=True, timeout=60
            )
            results.append({
                "action": "docker_prune",
                "output": result.stdout,
                "success": True,
            })
        except Exception as e:
            results.append({"action": "docker_prune", "success": False, "error": str(e)})
        
        # 清理旧日志
        try:
            result = subprocess.run(
                ["find", "/var/log", "-name", "*.log", "-mtime", "+7", "-delete"],
                capture_output=True, timeout=30
            )
            results.append({"action": "clean_old_logs", "success": True})
        except Exception as e:
            results.append({"action": "clean_old_logs", "success": False, "error": str(e)})
            
        return {"success": True, "actions": results}
    
    def execute(self, state: dict) -> dict:
        """根据诊断结果执行修复"""
        diagnosis = state.get("diagnosis", {})
        fixes = []
        
        for anomaly in state.get("anomalies", []):
            atype = anomaly["type"]
            if atype == "high_cpu":
                fixes.append(self.fix_high_cpu())
            elif atype == "high_memory":
                fixes.append(self.fix_high_memory())
            elif atype == "high_disk":
                fixes.append(self.fix_high_disk())
            elif atype == "container_crash":
                fixes.append(self.fix_container_crash())
        
        state["fix_actions"] = fixes
        state["fix_timestamp"] = datetime.now().isoformat()
        return state

# agents/verify.py
class VerifyAgent:
    def __init__(self, monitor: "MonitorAgent"):
        self.monitor = monitor
        
    def verify_metrics_improved(self, before: dict, after: dict) -> dict:
        """验证指标是否改善"""
        improvements = {}
        thresholds = self.monitor.config["alert_thresholds"]
        
        for key in ["cpu_percent", "memory_percent", "disk_percent"]:
            before_val = before.get(key, 0)
            after_val = after.get(key, 0)
            
            if before_val > thresholds.get(f"{key}_percent", 100):
                if after_val < before_val:
                    improvements[key] = {
                        "before": round(before_val, 2),
                        "after": round(after_val, 2),
                        "improved": True,
                    }
                else:
                    improvements[key] = {
                        "before": round(before_val, 2),
                        "after": round(after_val, 2),
                        "improved": False,
                    }
        
        all_fixed = all(v.get("improved") or v.get("before", 0) < thresholds.get(
            f"{v}" "_percent", 100) for v in improvements.values())
        
        return {
            "improvements": improvements,
            "all_fixed": all_fixed,
            "timestamp": after.get("timestamp", ""),
        }
    
    def check_service_health(self) -> dict:
        """检查关键服务是否恢复"""
        results = {}
        try:
            import docker
            client = docker.from_env()
            for container in client.containers.list():
                results[container.name] = {
                    "status": container.status,
                    "health": container.attrs.get("State", {}).get("Health", {}).get("Status", "unknown"),
                }
        except Exception as e:
            results["error"] = str(e)
        return results
    
    def execute(self, state: dict) -> dict:
        """执行验证"""
        # 重新采集指标
        new_monitor = MonitorAgent(self.monitor.config)
        new_metrics = new_monitor.collect_metrics()
        new_metrics["timestamp"] = new_monitor.config.get("last_check", "")
        
        verification = self.verify_metrics_improved(state["metrics"], new_metrics)
        service_health = self.check_service_health()
        
        state["verification"] = verification
        state["service_health"] = service_health
        state["metrics"] = new_metrics  # 更新为修复后的指标
        return state

# workflow.py
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator
import yaml
from datetime import datetime

class WorkflowState(TypedDict):
    timestamp: str
    metrics: dict
    anomalies: list
    severity: str
    diagnosis: dict
    fix_actions: list
    verification: dict
    service_health: dict
    report: str
    loop_count: int  # 防止无限循环

def create_workflow(config_path: str = "config.yaml"):
    with open(config_path) as f:
        config = yaml.safe_load(f)
    
    from agents.monitor import MonitorAgent
    from agents.diagnose import DiagnoseAgent
    from agents.fix import FixAgent
    from agents.verify import VerifyAgent
    from agents.report import ReportAgent
    from langchain_ollama import OllamaLLM
    
    llm = OllamaLLM(model=config["llm"]["model"], temperature=0)
    knowledge_base = []  # 从 SQLite 加载历史故障
    
    monitor = MonitorAgent(config)
    diagnose = DiagnoseAgent(llm, knowledge_base)
    fix = FixAgent(config)
    verify = VerifyAgent(monitor)
    report = ReportAgent(config, knowledge_base)
    
    # 定义节点
    def monitor_node(state: WorkflowState) -> WorkflowState:
        return monitor.execute(state)
    
    def diagnose_node(state: WorkflowState) -> WorkflowState:
        if not state["anomalies"]:
            return {**state, "diagnosis": {"root_cause": "无异常", "confidence": 1.0}}
        return diagnose.execute(state)
    
    def should_fix(state: WorkflowState) -> str:
        if not state["anomalies"]:
            return "report"
        if state["severity"] == "normal":
            return "report"
        if state["diagnosis"].get("requires_manual_review"):
            return "report"  # 需要人工确认
        return "fix"
    
    def fix_node(state: WorkflowState) -> WorkflowState:
        return fix.execute(state)
    
    def verify_node(state: WorkflowState) -> WorkflowState:
        return verify.execute(state)
    
    def should_loop(state: WorkflowState) -> str:
        if state.get("loop_count", 0) >= 3:
            return "report"
        if state.get("verification", {}).get("all_fixed"):
            return "report"
        return "diagnose"  # 修复后重新诊断
    
    def report_node(state: WorkflowState) -> WorkflowState:
        return report.execute(state)
    
    # 构建图
    graph = StateGraph(WorkflowState)
    
    graph.add_node("monitor", monitor_node)
    graph.add_node("diagnose", diagnose_node)
    graph.add_node("fix", fix_node)
    graph.add_node("verify", verify_node)
    graph.add_node("report", report_node)
    
    graph.set_entry_point("monitor")
    
    graph.add_conditional_edges(
        "monitor",
        lambda s: "diagnose" if s["anomalies"] else "report",
    )
    graph.add_conditional_edges(
        "diagnose",
        should_fix,
    )
    graph.add_edge("fix", "verify")
    graph.add_conditional_edges(
        "verify",
        should_loop,
    )
    graph.add_edge("report", END)
    
    app = graph.compile()
    return app

if __name__ == "__main__":
    workflow = create_workflow()
    
    initial_state: WorkflowState = {
        "timestamp": datetime.now().isoformat(),
        "metrics": {},
        "anomalies": [],
        "severity": "normal",
        "diagnosis": {},
        "fix_actions": [],
        "verification": {},
        "service_health": {},
        "report": "",
        "loop_count": 0,
    }
    
    result = workflow.invoke(initial_state)
    print(f"Workflow completed. Report: {result['report']}")

# agents/report.py
import requests
from datetime import datetime

class ReportAgent:
    def __init__(self, config: dict, knowledge_base: list):
        self.config = config
        self.kb = knowledge_base
        
    def generate_report(self, state: dict) -> str:
        """生成运维报告"""
        severity = state.get("severity", "normal")
        anomalies = state.get("anomalies", [])
        diagnosis = state.get("diagnosis", {})
        verification = state.get("verification", {})
        
        if not anomalies:
            return "✅ 系统健康检查完成，未发现异常。"
        
        lines = [f"🚨 VPS 故障自愈报告 ({severity.upper()})"]
        lines.append(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")
        
        lines.append("📊 异常详情:")
        for a in anomalies:
            lines.append(f"  • {a['message']} (当前: {a['value']:.1f})")
        
        lines.append("")
        lines.append("🧠 根因分析:")
        lines.append(f"  {diagnosis.get('root_cause', '未知')}")
        lines.append(f"  置信度: {diagnosis.get('confidence', 0):.0%}")
        
        lines.append("")
        lines.append("🛠️ 修复操作:")
        for fix in state.get("fix_actions", []):
            for action in fix.get("actions", []):
                status = "✅" if action.get("success") else "❌"
                lines.append(f"  {status} {action.get('action', 'unknown')}")
        
        lines.append("")
        lines.append("📈 修复效果:")
        verif = verification.get("improvements", {})
        for metric, data in verif.items():
            if data.get("improved"):
                lines.append(f"  ✅ {metric}: {data['before']:.1f}% → {data['after']:.1f}%")
            else:
                lines.append(f"  ⚠️  {metric}: 未改善")
        
        return "\n".join(lines)
    
    def send_notification(self, report: str):
        """发送通知"""
        provider = self.config.get("notification", {}).get("provider", "telegram")
        
        if provider == "telegram":
            token = self.config["notification"]["bot_token"]
            chat_id = self.config["notification"]["chat_id"]
            url = f"https://api.telegram.org/bot{token}/sendMessage"
            requests.post(url, json={
                "chat_id": chat_id,
                "text": report,
                "parse_mode": "HTML",
            }, timeout=10)
        elif provider == "email":
            # 邮件通知实现...
            pass
    
    def execute(self, state: dict) -> dict:
        report = self.generate_report(state)
        self.send_notification(report)
        state["report"] = report
        return state

# Dockerfile
FROM python:3.11-slim

RUN apt-get update && apt-get install -y \
    curl git procps \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

# 安装 Ollama（如需要本地运行）
# RUN curl -fsSL https://ollama.com/install.sh | sh

EXPOSE 8080
CMD ["python", "main.py"]

# docker-compose.yml
version: "3.8"
services:
  selfheal:
    build: .
    restart: unless-stopped
    environment:
      - TELEGRAM_BOT_TOKEN=${TELEGRAM_BOT_TOKEN}
      - TELEGRAM_CHAT_ID=${TELEGRAM_CHAT_ID}
      - OLLAMA_HOST=http://host.docker.internal:11434
    volumes:
      - ./knowledge:/app/knowledge
      - /var/run/docker.sock:/var/run/docker.sock
    depends_on:
      - prometheus
    networks:
      - ops-net

  prometheus:
    image: prom/prometheus:latest
    restart: unless-stopped
    volumes:
      - ./prometheus.yml:/etc/prometheus/prometheus.yml
      - prometheus-data:/prometheus
    networks:
      - ops-net

  node-exporter:
    image: prom/node-exporter:latest
    restart: unless-stopped
    network_mode: host
    pid: host
    networks:
      - ops-net

volumes:
  prometheus-data:

networks:
  ops-net:
    driver: bridge

# prometheus.yml
global:
  scrape_interval: 15s

scrape_configs:
  - job_name: 'node-exporter'
    static_configs:
      - targets: ['host.docker.internal:9100']

  - job_name: 'docker'
    docker_sd_configs:
      - host: unix:///var/run/docker.sock
        refresh_interval: 5s
    static_configs:
      - targets: ['localhost:9323']

# 使用 systemd timer 代替 cron（更可靠）
sudo tee /etc/systemd/system/vps-selfheal.service << 'EOF'
[Unit]
Description=VPS Self-Healing Agent
After=network.target

[Service]
Type=oneshot
User=root
WorkingDirectory=/root/vps-selfheal
ExecStart=/usr/bin/python3 main.py
EOF

sudo tee /etc/systemd/system/vps-selfheal.timer << 'EOF'
[Unit]
Description=Run VPS self-healing every 5 minutes

[Timer]
OnBootSec=2min
OnUnitActiveSec=5min
AccuracySec=1s

[Install]
WantedBy=timers.target
EOF

sudo systemctl enable vps-selfheal.timer
sudo systemctl start vps-selfheal.timer

# 查看状态
systemctl status vps-selfheal.timer