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

# 拉取适合运维场景的模型（参数少、响应快）
ollama pull qwen2.5:7b
ollama pull nomic-embed-text  # 用于知识库向量嵌入

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:
    """执行系统命令，返回结果和状态"""
    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": cmd}
    except Exception as e:
        return {"success": False, "error": str(e), "command": cmd}

def get_system_status() -> dict:
    """获取系统整体状态"""
    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:
    """获取占用资源最多的进程"""
    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:
    """检查 systemd 服务状态"""
    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:
    """获取 Docker 容器状态"""
    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_core.runnables import RunnablePassthrough
from langchain_ollama import ChatOllama
from tools.system_tools import (
    run_command, get_system_status, get_top_processes,
    check_service, docker_status
)

# 定义运维工具
@tool
def system_info() -> str:
    """获取 VPS 系统整体状态：CPU、内存、磁盘、负载"""
    status = get_system_status()
    return json.dumps(status, ensure_ascii=False, indent=2)

@tool
def top_processes(sort_by: str = "memory", n: int = 10) -> str:
    """获取占用资源最多的进程列表，sort_by 可选 memory 或 cpu"""
    procs = get_top_processes(n=n, sort_by=sort_by)
    return json.dumps(procs, ensure_ascii=False, indent=2)

@tool
def exec_command(command: str) -> str:
    """在 VPS 上执行命令（需管理员确认）"""
    result = run_command(command)
    return json.dumps(result, ensure_ascii=False, indent=2)

@tool
def check_service_status(service: str) -> str:
    """检查指定 systemd 服务的运行状态"""
    return json.dumps(check_service(service), ensure_ascii=False)

@tool
def docker_list() -> str:
    """列出所有 Docker 容器及其状态"""
    return json.dumps(docker_status(), ensure_ascii=False)

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

from langchain.agents import create_tool_calling_agent, AgentExecutor

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

# 构建提示词——加入安全约束
SYSTEM_PROMPT = """你是一个专业的 VPS 运维 AI 助手（Ops Copilot）。
你的职责是帮助用户诊断问题、执行操作、提供建议。

## 行为准则
1. **安全第一**：执行破坏性操作（删除文件、重启服务、停止容器）前必须向用户确认
2. **精准诊断**：先用只读工具收集信息，再决定是否需要执行修改操作
3. **解释清晰**：用中文回答，解释技术的含义和影响
4. **承认未知**：不确定时明确说明，不要编造答案
5. **拒绝危险**：拒绝执行明显有害的操作（如 rm -rf /、格式化磁盘等）

## 当前时间
{current_time}

## 可用工具
{tools}

现在，请开始帮助用户。""".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"
)

# 完整的 Copilot 链路
copilot = {
    "input": lambda x: x["input"],
    "chat_history": lambda x: memory.load_memory_variables(x)["chat_history"],
    "agent_scratchpad": lambda x: format_agent_scratchpad(x["intermediate_steps"]),
} | agent_executor

# 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):
        """从知识库目录构建向量索引"""
        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:
        """检索相关知识"""
        if not self.vectorstore:
            return []
        return self.vectorstore.similarity_search(query, k=k)

# 在 agent prompt 中加入知识检索
KB_PROMPT = """
## 运维知识库
以下知识来自你的历史运维记录和操作手册，请参考这些信息来回答用户问题：

{context}

如果知识库中没有相关信息，请基于你的通用知识回答，并说明这是通用建议。
"""

# 构建 RAG 增强链
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain

kb = OpsKnowledgeBase()
kb.build()

retriever = kb.vectorstore.as_retriever(search_kwargs={"k": 3})

doc_chain = create_stuff_documents_chain(llm, KB_PROMPT)
rag_chain = create_retrieval_chain(retriever, doc_chain)

# 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'\bswapoff\s+-a\b',       # 关闭全部 swap
    r'\breboot\b',             # 重启系统
    r'\ ipython\b',             # 交互式 shell
]

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:
    """检查命令是否安全"""
    for pattern in DANGEROUS_PATTERNS:
        if re.search(pattern, cmd, re.IGNORECASE):
            return {"safe": False, "reason": f"匹配危险模式: {pattern}"}
    
    base_cmd = cmd.strip().split()[0] if cmd.strip() else ""
    if base_cmd not in ALLOWED_COMMANDS:
        return {"safe": None, "reason": f"未在白名单中的命令: {base_cmd}，需要人工确认"}
    
    return {"safe": True, "reason": "命令在白名单中"}

def safe_execute(agent_executor, input_text: str, history: list = None) -> dict:
    """带安全控制的执行入口"""
    safety = is_safe_command(input_text)
    
    if safety["safe"] is False:
        return {
            "status": "blocked",
            "message": f"⛔ 操作被拦截：{safety['reason']}",
            "suggestion": "请联系管理员审核此操作"
        }
    
    if safety["safe"] is None:
        return {
            "status": "pending_confirm",
            "message": f"⚠️ 检测到未白名单命令，请确认是否执行：\n\n\`{input_text}\`\n\n回复 '确认' 执行，或 '取消' 跳过。"
        }
    
    # 正常执行
    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/*

# 安装 Ollama
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 已就绪！\n\n"
        "直接发送问题即可，例如：\n"
        "• '系统状态如何？'\n"
        "• '哪些进程占用 CPU 最高？'\n"
        "• '重启 nginx 服务'\n"
        "• '帮我查看今天的错误日志'"
    )

@bot.message_handler(func=lambda m: True)
def handle_message(message):
    chat_id = message.chat.id
    user_input = message.text.strip()
    
    with bot.state_privacy(chat_id, True):
        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
import threading

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:]  # 保留最近 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):
    """将故障排查过程保存为知识"""
    doc = f"""## 故障案例
**问题**: {failure_scenario}
**解决方案**: {solution}
**时间**: {datetime.now().isoformat()}
"""
    with open(f"./knowledge/incidents/{datetime.now().strftime('%Y%m%d')}.md", "w") as f:
        f.write(doc)

def refine_promptBasedOnFeedback(old_prompt: str, feedback: str) -> str:
    """用 LLM 自身优化自己的提示词"""
    optimizer_llm = ChatOllama(model="qwen2.5:7b", temperature=0.3)
    
    improvement_request = f"""
    以下是一个运维 Copilot 的系统提示词：
    ---
    {old_prompt}
    ---
    
    用户反馈：{feedback}
    
    请改进提示词，使 Copilot 能更好地处理此类问题。只输出改进后的提示词，不要其他内容。
    """
    
    return optimizer_llm.invoke(improvement_request).content