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

# Pull models suitable for ops scenarios (lightweight, fast)
ollama pull qwen2.5:7b
ollama pull nomic-embed-text  # for knowledge base embeddings

pip install langchain langchain-community langchain-core \
            chromadb ollama python-dotenv \
            psutil pydantic

# tools/system_tools.py
import subprocess
import psutil
import json
from datetime import datetime

def run_command(cmd: str, timeout: int = 30) -> dict:
    """Execute system command, return result and status"""
    try:
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True, timeout=timeout
        )
        return {
            "success": result.returncode == 0,
            "stdout": result.stdout.strip(),
            "stderr": result.stderr.strip(),
            "return_code": result.returncode,
            "timestamp": datetime.now().isoformat()
        }
    except subprocess.TimeoutExpired:
        return {"success": False, "error": "Command timeout", "command": cmd}
    except Exception as e:
        return {"success": False, "error": str(e), "command": cmd}

def get_system_status() -> dict:
    """Get overall system status"""
    return {
        "cpu_percent": psutil.cpu_percent(interval=1),
        "memory": {
            "total_gb": round(psutil.virtual_memory().total / 1024**3, 2),
            "available_gb": round(psutil.virtual_memory().available / 1024**3, 2),
            "used_percent": psutil.virtual_memory().percent
        },
        "disk": {
            "/": {
                "total_gb": round(psutil.disk_usage('/').total / 1024**3, 2),
                "used_gb": round(psutil.disk_usage('/').used / 1024**3, 2),
                "used_percent": psutil.disk_usage('/').percent
            }
        },
        "load_avg": psutil.getloadavg(),
        "uptime_seconds": psutil.boot_time()
    }

def get_top_processes(n: int = 10, sort_by: str = "memory") -> list:
    """Get top resource-consuming processes"""
    processes = []
    for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent', 'status']):
        try:
            processes.append(proc.info)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    
    key_map = {"cpu": "cpu_percent", "memory": "memory_percent"}
    processes.sort(
        key=lambda x: x.get(key_map.get(sort_by, "memory_percent"), 0), 
        reverse=True
    )
    return processes[:n]

def check_service(service_name: str) -> dict:
    """Check systemd service status"""
    result = run_command(f"systemctl is-active {service_name}")
    return {
        "service": service_name,
        "status": result["stdout"],
        "enabled": run_command(f"systemctl is-enabled {service_name}")["stdout"]
    }

def docker_status() -> dict:
    """Get Docker container status"""
    result = run_command("docker ps --format '{{.Names}}\t{{.Status}}\t{{.Ports}}'")
    containers = []
    for line in result["stdout"].split("\n"):
        if line.strip():
            parts = line.strip().split("\t")
            containers.append({
                "name": parts[0] if len(parts) > 0 else "",
                "status": parts[1] if len(parts) > 1 else "",
                "ports": parts[2] if len(parts) > 2 else ""
            })
    return {"containers": containers, "count": len(containers)}

# copilot/core.py
from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_ollama import ChatOllama
from tools.system_tools import (
    run_command, get_system_status, get_top_processes,
    check_service, docker_status
)
import json

@tool
def system_info() -> str:
    """Get VPS system overall status: CPU, memory, disk, load"""
    status = get_system_status()
    return json.dumps(status, indent=2)

@tool
def top_processes(sort_by: str = "memory", n: int = 10) -> str:
    """Get top resource-consuming process list, sort_by: memory or cpu"""
    procs = get_top_processes(n=n, sort_by=sort_by)
    return json.dumps(procs, indent=2)

@tool
def exec_command(command: str) -> str:
    """Execute command on VPS (requires admin confirmation)"""
    result = run_command(command)
    return json.dumps(result, indent=2)

@tool
def check_service_status(service: str) -> str:
    """Check specified systemd service status"""
    return json.dumps(check_service(service))

@tool
def docker_list() -> str:
    """List all Docker containers and their status"""
    return json.dumps(docker_status())

TOOLS = [system_info, top_processes, exec_command, check_service_status, docker_list]

from langchain.agents import create_tool_calling_agent, AgentExecutor

# Initialize LLM
llm = ChatOllama(
    model="qwen2.5:7b",
    temperature=0,
    base_url="http://localhost:11434"
)

# Build prompt with safety constraints
SYSTEM_PROMPT = """You are a professional VPS operations AI assistant (Ops Copilot).
Your duty is to help users diagnose issues, execute operations, and provide recommendations.

## Rules
1. **Safety first**: Always confirm with the user before executing destructive operations 
   (delete files, restart services, stop containers)
2. **Precise diagnosis**: Use read-only tools to gather info first, then decide if 
   modification is needed
3. **Clear explanation**: Explain technical meaning and impact in clear language
4. **Admit uncertainty**: Clearly state when you're unsure, don't fabricate answers
5. **Refuse dangerous requests**: Reject obviously harmful operations (rm -rf /, 
   disk formatting, etc.)

## Available Tools
{tools}

Now, help the user. """.strip()

prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    MessagesPlaceholder("chat_history", optional=True),
    ("human", "{input}"),
    MessagesPlaceholder("agent_scratchpad"),
])

agent = create_tool_calling_agent(llm, TOOLS, prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=TOOLS,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=15,
    return_intermediate_steps=True
)

from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True,
    output_key="output"
)

# knowledge_base.py
import chromadb
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter

class OpsKnowledgeBase:
    def __init__(self, docs_dir: str = "./knowledge"):
        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
        self.client = chromadb.PersistentClient(path="./chroma_ops_db")
        self.docs_dir = docs_dir
        self.vectorstore = None
    
    def build(self):
        """Build vector index from knowledge base directory"""
        loader = DirectoryLoader(
            self.docs_dir,
            glob="**/*.md",
            loader_cls=TextLoader,
            loader_kwargs={"encoding": "utf-8"}
        )
        documents = loader.load()
        
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            separators=["\n\n", "\n", ".", " "]
        )
        chunks = splitter.split_documents(documents)
        
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            client=self.client,
            collection_name="ops_knowledge"
        )
        return len(chunks)
    
    def search(self, query: str, k: int = 3) -> list:
        """Search relevant knowledge"""
        if not self.vectorstore:
            return []
        return self.vectorstore.similarity_search(query, k=k)

KB_PROMPT = """
## Ops Knowledge Base
The following knowledge comes from your historical ops records and operation manuals.
Please reference this information when answering user questions:

{context}

If the knowledge base doesn't have relevant info, answer based on your general 
knowledge and note that this is a general recommendation.
"""

# security.py
import re

DANGEROUS_PATTERNS = [
    r'\brm\s+-rf\b',
    r'\bformat\b.*\b/',
    r'\bmkfs\b',
    r'\bdd\s+if=',
    r'\bwget\s+.*\|.*sh\b',
    r'\bcurl\s+.*\|.*sh\b',
    r'\bchown\s+-R\s+root\b',
    r'\biptables\s+-F\b',
    r'\breboot\b',
]

ALLOWED_COMMANDS = [
    'ls', 'cat', 'head', 'tail', 'grep', 'find', 'du', 'df',
    'top', 'htop', 'ps', 'free', 'uptime', 'who', 'uname',
    'docker', 'systemctl', 'journalctl', 'ss', 'netstat',
    'curl', 'wget', 'ping', 'dig', 'nslookup', 'traceroute',
    'apt', 'apt-get', 'yum', 'dnf', 'pip', 'npm',
    'git', 'tar', 'gzip', 'zip', 'unzip',
]

def is_safe_command(cmd: str) -> dict:
    """Check if command is safe"""
    for pattern in DANGEROUS_PATTERNS:
        if re.search(pattern, cmd, re.IGNORECASE):
            return {"safe": False, "reason": f"Matched danger pattern: {pattern}"}
    
    base_cmd = cmd.strip().split()[0] if cmd.strip() else ""
    if base_cmd not in ALLOWED_COMMANDS:
        return {"safe": None, "reason": f"Command not in whitelist: {base_cmd}, needs human confirmation"}
    
    return {"safe": True, "reason": "Command is in whitelist"}

def safe_execute(agent_executor, input_text: str, history: list = None) -> dict:
    """Execution entry with safety controls"""
    safety = is_safe_command(input_text)
    
    if safety["safe"] is False:
        return {
            "status": "blocked",
            "message": f"⛔ Operation blocked: {safety['reason']}",
            "suggestion": "Please contact administrator for review"
        }
    
    if safety["safe"] is None:
        return {
            "status": "pending_confirm",
            "message": f"⚠️ Non-whitelisted command detected. Confirm execution?\n\n`{input_text}`\n\nReply 'confirm' to proceed or 'cancel' to skip."
        }
    
    result = agent_executor.invoke({
        "input": input_text,
        "chat_history": history or []
    })
    return {"status": "success", "output": result["output"]}

# Dockerfile
FROM python:3.11-slim

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

RUN curl -fsSL https://ollama.com/install.sh | sh

WORKDIR /app
COPY . .

RUN pip install --no-cache-dir -r requirements.txt

RUN ollama pull qwen2.5:7b && \
    ollama pull nomic-embed-text

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

# bots/telegram_bot.py
import telebot
from copilot.core import copilot

BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
bot = telebot.TeleBot(BOT_TOKEN)

@bot.message_handler(commands=['start', 'help'])
def send_welcome(message):
    bot.reply_to(message, 
        "🤖 Ops Copilot is ready!\n\n"
        "Send questions directly, e.g.:\n"
        "• 'How is system status?'\n"
        "• 'Which processes use the most CPU?'\n"
        "• 'Restart nginx service'\n"
        "• 'Show me today's error logs'"
    )

@bot.message_handler(func=lambda m: True)
def handle_message(message):
    chat_id = message.chat.id
    user_input = message.text.strip()
    
    result = safe_execute(copilot, user_input)
    
    if result["status"] == "blocked":
        bot.reply_to(message, result["message"])
    elif result["status"] == "pending_confirm":
        confirm_states[chat_id] = {"result": result, "input": user_input}
        bot.reply_to(message, result["message"])
    else:
        bot.reply_to(message, result["output"])

# main.py
from flask import Flask, request, jsonify
from copilot.core import copilot

app = Flask(__name__)
chat_states = {}

@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    chat_id = data.get('chat_id', 'default')
    message = data.get('message', '')
    
    history = chat_states.get(chat_id, [])
    result = safe_execute(copilot, message, history)
    
    if result["status"] == "success":
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": result["output"]})
        chat_states[chat_id] = history[-20:]
    
    return jsonify(result)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080)

def save_incident_to_kb(failure_scenario: str, solution: str):
    """Save troubleshooting process as knowledge"""
    doc = f"""## Incident Case
**Issue**: {failure_scenario}
**Solution**: {solution}
**Time**: {datetime.now().isoformat()}
"""
    import os
    os.makedirs(f"./knowledge/incidents/", exist_ok=True)
    with open(f"./knowledge/incidents/{datetime.now().strftime('%Y%m%d')}.md", "w") as f:
        f.write(doc)

def refine_prompt_based_on_feedback(old_prompt: str, feedback: str) -> str:
    """Use LLM to optimize its own prompt"""
    optimizer_llm = ChatOllama(model="qwen2.5:7b", temperature=0.3)
    
    improvement_request = f"""
    Current system prompt:
    ---
    {old_prompt}
    ---
    
    User feedback: {feedback}
    
    Improve the prompt so the Copilot handles such issues better. 
    Output only the improved prompt, nothing else.
    """
    
    return optimizer_llm.invoke(improvement_request).content