# docker-compose.agent.yml
services:
  # ── 数据采集层 ──
  node-exporter:
    image: prom/node-exporter:latest
    # 采集所有 VPS 的系统指标

  promtail:
    image: grafana/promtail:latest
    # 收集所有 VPS 的日志并推送至 Loki

  # ── 数据存储层 ──
  prometheus:
    image: prom/prometheus:latest
    # 存储时序指标

  loki:
    image: grafana/loki:latest
    # 存储日志

  grafana:
    image: grafana/grafana:latest
    # 可视化 + 告警规则

  # ── AI Agent 核心层 ──
  agent-orchestrator:
    build: ./agent
    # Agent 编排器 —— 系统大脑
    environment:
      - LLM_ENDPOINT=http://ollama:11434
      - KNOWLEDGE_BASE_PATH=/data/knowledge-base
      - DANGER_COMMANDS=rm -rf,shutdown,reboot
      - APPROVAL_REQUIRED=true

  ollama:
    image: ollama/ollama:latest
    # 本地 LLM 推理引擎
    volumes:
      - ./models:/models

  # ── 执行层 ──
  agent-executor:
    build: ./executor
    # 在远程 VPS 上运行的轻量 Agent
    # 接收指令并安全执行

# agent/orchestrator.py
import asyncio
import json
from datetime import datetime
from pathlib import Path
from typing import Optional

from prometheus_api_client import PrometheusConnect
from langchain_community.llms.ollama import Ollama
from langchain.prompts import PromptTemplate

class VPSAgentOrchestrator:
    def __init__(self, llm_model="qwen2.5:7b-instruct"):
        self.prom = PrometheusConnect(url="http://prometheus:9090")
        self.llm = Ollama(model=llm_model, base_url="http://ollama:11434")
        self.knowledge_base = Path("/data/knowledge-base")
        self.dangerous_commands = {"rm -rf", "shutdown", "reboot", "dd"}
        self.approval_required = True

    async def run_cycle(self):
        """主循环：每个周期执行一次完整感知-推理-决策-执行"""
        print(f"[{datetime.now()}] Running agent cycle...")

        # ① 感知：采集当前状态
        status = await self.perceive()

        # ② 推理：分析是否有异常
        analysis = await self.reason(status)

        if not analysis["has_anomaly"]:
            print("  No anomalies detected.")
            return

        # ③ 决策：生成修复方案
        plan = await self.decide(analysis)

        # ④ 执行：在安全边界内执行
        await self.execute(plan)

    async def perceive(self) -> dict:
        """从多源采集数据"""
        return {
            "metrics": self._fetch_metrics(),
            "logs": self._fetch_recent_logs(),
            "events": self._fetch_alert_events(),
            "timestamp": datetime.now().isoformat(),
        }

    async def reason(self, status: dict) -> dict:
        """LLM 分析异常并定位根因"""
        prompt = PromptTemplate.from_template("""
You are a senior SRE engineer. Analyze the following VPS status data
and determine if there are anomalies, what the root cause might be,
and what action should be taken.

Status Data:
{status}

Respond in JSON format:
{{
  "has_anomaly": true/false,
  "severity": "critical/high/medium/low",
  "root_cause": "description of root cause",
  "evidence": ["list of supporting evidence"],
  "recommended_action": "description of recommended action",
  "action_command": "specific command to execute (or null if manual review needed)"
}}
""")

        chain = prompt | self.llm
        response = await chain.ainvoke({"status": json.dumps(status, indent=2)})
        return json.loads(response)

    async def decide(self, analysis: dict) -> dict:
        """生成可执行的修复方案"""
        risk_score = self._calculate_risk(analysis.get("action_command", ""))

        plan = {
            "analysis": analysis,
            "risk_score": risk_score,
            "requires_approval": risk_score >= 7 or self.approval_required,
            "executed_at": None,
        }

        if plan["requires_approval"]:
            await self._send_approval_request(plan)

        return plan

    def _calculate_risk(self, command: str) -> int:
        """计算操作风险等级 (0-10)"""
        if not command:
            return 0
        score = 0
        for dangerous in self.dangerous_commands:
            if dangerous in command:
                score += 5
        if "systemctl restart" in command:
            score += 2
        if "iptables" in command or "firewall" in command:
            score += 3
        return min(score, 10)

    async def execute(self, plan: dict):
        """执行修复方案"""
        if plan["requires_approval"]:
            print("  Action requires manual approval. Sent notification.")
            return

        command = plan["analysis"].get("action_command")
        if command:
            print(f"  Executing: {command}")
            # 在远程 VPS 上安全执行
            result = await self._safe_execute(command)
            plan["executed_at"] = datetime.now().isoformat()
            plan["result"] = result
            await self._record_learning(plan)

    async def _record_learning(self, plan: dict):
        """将本次事件记录到知识库，供未来参考"""
        record = {
            "timestamp": plan["executed_at"],
            "analysis": plan["analysis"],
            "risk_score": plan["risk_score"],
            "outcome": plan.get("result", "pending"),
        }
        log_file = self.knowledge_base / "events" / f"{datetime.now().strftime('%Y%m%d')}.jsonl"
        log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(log_file, "a") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

# executor/agent.py
"""
在远程 VPS 上运行的轻量级 Agent，接收指令并安全执行。
"""
import asyncio
import json
import subprocess
from pathlib import Path
from datetime import datetime

class SafeExecutor:
    DANGEROUS_PATTERNS = [
        "rm -rf /", "mkfs", "dd if=", "chmod 777",
        "> /dev/sda", "wget .* | sh", "curl .* | bash"
    ]

    def __init__(self, max_output_size=10240):
        self.max_output_size = max_output_size
        self.audit_log = Path("/var/log/agent-executor/audit.log")

    async def execute(self, command: str, timeout: int = 30) -> dict:
        """安全执行命令"""
        # 安全检查
        if self._is_dangerous(command):
            return {"error": "Command blocked by safety policy", "command": command}

        # 记录审计日志
        self._audit_log(command)

        try:
            result = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                limit=self.max_output_size,
            )
            stdout, stderr = await asyncio.wait_for(result.communicate(), timeout=timeout)

            return {
                "command": command,
                "exit_code": result.returncode,
                "stdout": stdout.decode()[:self.max_output_size],
                "stderr": stderr.decode()[:self.max_output_size],
                "success": result.returncode == 0,
            }
        except asyncio.TimeoutError:
            result.kill()
            return {"error": "Command timed out", "command": command}

    def _is_dangerous(self, command: str) -> bool:
        for pattern in self.DANGEROUS_PATTERNS:
            if pattern in command:
                return True
        return False

    def _audit_log(self, command: str):
        entry = {
            "timestamp": datetime.now().isoformat(),
            "command": command,
            "action": "executing",
        }
        self.audit_log.parent.mkdir(parents=True, exist_ok=True)
        with open(self.audit_log, "a") as f:
            f.write(json.dumps(entry) + "\n")

# 在一台独立的 VPS 上部署 Agent 中枢（推荐 4C8G）
git clone https://github.com/selfvps/ai-vps-agent.git
cd ai-vps-agent

# 安装依赖
pip install -r requirements.txt

# 配置 Ollama（本地 LLM）
docker compose up -d ollama prometheus loki grafana
ollama pull qwen2.5:7b-instruct

# 在每台需要管理的 VPS 上运行
curl -sSL https://raw.githubusercontent.com/selfvps/ai-vps-agent/main/install.sh | bash

# 配置连接到 Agent 中枢
cat > /etc/agent/config.yaml <<EOF
orchestrator:
  host: <your-agent-server-ip>
  port: 8080
  
safety:
  approval_required: true
  dangerous_commands:
    - "rm -rf /"
    - "shutdown"
    - "reboot"
EOF

# 启动 Agent
systemctl enable --now vps-agent

# 查看 Agent 状态
journalctl -u vps-agent -f

# 模拟一个问题，验证自动响应
# （例如：临时创建一个大文件模拟磁盘满）
sudo dd if=/dev/zero of=/tmp/fill_disk bs=1M count=500
# 观察 Agent 是否自动检测并清理