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

# Pull lightweight models (suitable for VPS resources)
ollama pull llama3.2:3b
ollama pull nomic-embed-text  # for RAG retrieval

# Verify installation
ollama list

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

# Create project structure
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  # seconds
  
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:
        """Query Prometheus metrics"""
        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:
        """Collect key system metrics"""
        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:
        """Detect anomalies"""
        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"High CPU usage: {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"High memory usage: {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"High disk usage: {metrics['disk_percent']:.1f}%"
            })
        
        if metrics["container_restarts"] > thresholds["container_restart_count"]:
            anomalies.append({
                "type": "container_crash",
                "value": metrics["container_restarts"],
                "threshold": thresholds["container_restart_count"],
                "message": f"Container frequent restarts: {metrics['container_restart_count']:.0f} times/10min"
            })
        
        return anomalies
    
    def determine_severity(self, anomalies: list) -> str:
        """Determine alert severity"""
        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 execution logic"""
        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  # Historical fault knowledge base
        
    def build_prompt(self, anomalies: list, metrics: dict) -> str:
        """Build diagnostic prompt"""
        kb_context = "\n".join([
            f"- {item['incident']}: {item['root_cause']} → {item['solution']}"
            for item in self.kb[-5:]  # Last 5 entries
        ])
        
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a VPS ops expert. Analyze root cause based on monitoring anomalies and system metrics.
Historical fault cases reference:
{kb_context}

Output JSON format:
{{
  "root_cause": "Root cause description",
  "confidence": 0.0-1.0,
  "suggested_fixes": ["Fix plan 1", "Fix plan 2"],
  "risk_level": "low|medium|high",
  "requires_manual_review": true|false
}}"""),
            ("user", """Current anomalies: {anomalies}
Current metrics: {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:
        """Execute diagnosis"""
        prompt = self.build_prompt(state["anomalies"], state["metrics"])
        response = self.llm.invoke(prompt)
        
        try:
            diagnosis = json.loads(response.content)
        except json.JSONDecodeError:
            diagnosis = {
                "root_cause": "Unable to auto-diagnose, manual intervention required",
                "confidence": 0.0,
                "suggested_fixes": ["Contact ops personnel"],
                "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:
        """Handle high CPU anomaly"""
        results = []
        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:
                results.append({
                    "action": "log_alert",
                    "message": f"High CPU process: 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:
        """Handle high memory anomaly"""
        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,
                    "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:
        """Handle frequent container restarts"""
        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"],
                    })
            
            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:
        """Handle insufficient disk space"""
        results = []
        
        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:
            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:
        """Execute fixes based on 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):
        self.monitor = monitor
        
    def verify_metrics_improved(self, before: dict, after: dict) -> dict:
        """Verify metrics improvement"""
        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(key, 100)
            for key, v in improvements.items()
        )
        
        return {
            "improvements": improvements,
            "all_fixed": all_fixed,
            "timestamp": after.get("timestamp", ""),
        }
    
    def check_service_health(self) -> dict:
        """Check if key services recovered"""
        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:
        """Execute verification"""
        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
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 = []
    
    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": "No anomaly", "confidence": 1.0}}
        return diagnose.execute(state)
    
    def should_fix(state: WorkflowState) -> str:
        if not state["anomalies"] or 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)
    
    return graph.compile()

# 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:
        """Generate ops report"""
        severity = state.get("severity", "normal")
        anomalies = state.get("anomalies", [])
        diagnosis = state.get("diagnosis", {})
        verification = state.get("verification", {})
        
        if not anomalies:
            return "✅ System health check complete, no anomalies found."
        
        lines = [f"🚨 VPS Self-Healing Report ({severity.upper()})"]
        lines.append(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")
        
        lines.append("📊 Anomaly Details:")
        for a in anomalies:
            lines.append(f"  • {a['message']} (current: {a['value']:.1f})")
        
        lines.append("")
        lines.append("🧠 Root Cause Analysis:")
        lines.append(f"  {diagnosis.get('root_cause', 'Unknown')}")
        lines.append(f"  Confidence: {diagnosis.get('confidence', 0):.0%}")
        
        lines.append("")
        lines.append("🛠️ Fix Actions:")
        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("📈 Fix Results:")
        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}: Not improved")
        
        return "\n".join(lines)
    
    def send_notification(self, report: str):
        """Send notification"""
        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)
    
    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 . .

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

# Use systemd timer (more reliable than 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