ai-status-light/internal/api/api.go

package api

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"AI-Status-Light/internal/database"
)

type Server struct {
	db        *database.DB
	server    *http.Server
	clients   map[*websocket.Conn]bool
	clientsMu sync.Mutex
	upgrader  websocket.Upgrader
}

type Response struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

func New(db *database.DB, addr string) *Server {
	s := &Server{
		db:      db,
		clients: make(map[*websocket.Conn]bool),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/api/mqtt", s.handleMQTT)
	mux.HandleFunc("/api/mqtt/", s.handleMQTTByID)
	mux.HandleFunc("/api/health", s.handleHealth)
	mux.HandleFunc("/ws", s.handleWebSocket)
	mux.HandleFunc("/", s.handleIndex)

	s.server = &http.Server{
		Addr:    addr,
		Handler: corsMiddleware(mux),
	}
	return s
}

func (s *Server) Start() error {
	return s.server.ListenAndServe()
}

func corsMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	})
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, Response{Code: 0, Message: "ok"})
}

func (s *Server) handleMQTT(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		s.listMQTTConfigs(w, r)
	case http.MethodPost:
		s.createMQTTConfig(w, r)
	default:
		writeJSON(w, http.StatusMethodNotAllowed, Response{Code: -1, Message: "方法不允许"})
	}
}

func (s *Server) handleMQTTByID(w http.ResponseWriter, r *http.Request) {
	idStr := strings.TrimPrefix(r.URL.Path, "/api/mqtt/")
	id, err := strconv.Atoi(idStr)
	if err != nil {
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "无效的 ID"})
		return
	}

	switch r.Method {
	case http.MethodGet:
		s.getMQTTConfig(w, id)
	case http.MethodPut:
		s.updateMQTTConfig(w, r, id)
	case http.MethodDelete:
		s.deleteMQTTConfig(w, id)
	default:
		writeJSON(w, http.StatusMethodNotAllowed, Response{Code: -1, Message: "方法不允许"})
	}
}

func (s *Server) listMQTTConfigs(w http.ResponseWriter, r *http.Request) {
	configs, err := s.db.ListMQTTConfigs()
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}
	writeJSON(w, http.StatusOK, Response{Code: 0, Message: "ok", Data: configs})
}

func (s *Server) createMQTTConfig(w http.ResponseWriter, r *http.Request) {
	var cfg database.MQTTConfig
	if err := json.NewDecoder(r.Body).Decode(&cfg); err != nil {
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "无效的请求体"})
		return
	}

	if cfg.Broker == "" {
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "broker 不能为空"})
		return
	}

	if cfg.ClientID == "" {
		cfg.ClientID = "opencode-monitor"
	}
	if cfg.Topic == "" {
		cfg.Topic = "opencode/status"
	}

	if err := s.db.SaveMQTTConfig(&cfg); err != nil {
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	writeJSON(w, http.StatusCreated, Response{Code: 0, Message: "创建成功", Data: cfg})
}

func (s *Server) getMQTTConfig(w http.ResponseWriter, id int) {
	configs, err := s.db.ListMQTTConfigs()
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	for _, cfg := range configs {
		if cfg.ID == id {
			writeJSON(w, http.StatusOK, Response{Code: 0, Message: "ok", Data: cfg})
			return
		}
	}

	writeJSON(w, http.StatusNotFound, Response{Code: -1, Message: "配置不存在"})
}

func (s *Server) updateMQTTConfig(w http.ResponseWriter, r *http.Request, id int) {
	var cfg database.MQTTConfig
	if err := json.NewDecoder(r.Body).Decode(&cfg); err != nil {
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "无效的请求体"})
		return
	}

	cfg.ID = id
	if cfg.Broker == "" {
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "broker 不能为空"})
		return
	}

	if cfg.ClientID == "" {
		cfg.ClientID = "opencode-monitor"
	}
	if cfg.Topic == "" {
		cfg.Topic = "opencode/status"
	}

	if err := s.db.SaveMQTTConfig(&cfg); err != nil {
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	writeJSON(w, http.StatusOK, Response{Code: 0, Message: "更新成功", Data: cfg})
}

func (s *Server) deleteMQTTConfig(w http.ResponseWriter, id int) {
	if err := s.db.DeleteMQTTConfig(id); err != nil {
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}
	writeJSON(w, http.StatusOK, Response{Code: 0, Message: "删除成功"})
}

func writeJSON(w http.ResponseWriter, statusCode int, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(statusCode)
	json.NewEncoder(w).Encode(data)
}

func (s *Server) GetAddr() string {
	return s.server.Addr
}

func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}

	s.clientsMu.Lock()
	s.clients[conn] = true
	s.clientsMu.Unlock()

	log.Printf("WebSocket 客户端已连接，当前连接数: %d", len(s.clients))

	go func() {
		defer func() {
			s.clientsMu.Lock()
			delete(s.clients, conn)
			s.clientsMu.Unlock()
			conn.Close()
			log.Printf("WebSocket 客户端已断开，当前连接数: %d", len(s.clients))
		}()

		for {
			_, _, err := conn.ReadMessage()
			if err != nil {
				break
			}
		}
	}()
}

func (s *Server) BroadcastStatus(port int, status string, code string) {
	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()

	if len(s.clients) == 0 {
		return
	}

	payload := map[string]interface{}{
		"port":      port,
		"status":    status,
		"code":      code,
		"timestamp": time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return
	}

	for client := range s.clients {
		err := client.WriteMessage(websocket.TextMessage, data)
		if err != nil {
			client.Close()
			delete(s.clients, client)
		}
	}
}

func (s *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	html := `<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>OpenCode Monitor</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: #f5f5f5; padding: 20px; }
        .container { max-width: 1200px; margin: 0 auto; }
        h1 { color: #333; margin-bottom: 20px; }
        .status-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px; }
        .status-card { background: white; border-radius: 12px; padding: 24px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); }
        .status-card h2 { color: #666; font-size: 14px; margin-bottom: 8px; }
        .status-value { font-size: 32px; font-weight: bold; margin-bottom: 8px; }
        .status-time { color: #999; font-size: 12px; }
        .status-空闲 { color: #52c41a; }
        .status-忙碌 { color: #ff4d4f; }
        .status-思考中 { color: #faad14; }
        .status-运行中 { color: #1890ff; }
        .status-完成 { color: #52c41a; }
        .status-错误 { color: #ff4d4f; }
        .status-重试中 { color: #faad14; }
        .status-修改中 { color: #722ed1; }
        .log { margin-top: 20px; background: white; border-radius: 12px; padding: 24px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); }
        .log h2 { color: #666; font-size: 14px; margin-bottom: 12px; }
        .log-list { max-height: 300px; overflow-y: auto; }
        .log-item { padding: 8px 0; border-bottom: 1px solid #f0f0f0; font-size: 14px; }
        .log-item:last-child { border-bottom: none; }
        .log-time { color: #999; margin-right: 8px; }
        .log-port { color: #1890ff; margin-right: 8px; }
        .connected { color: #52c41a; font-size: 12px; margin-left: 10px; }
        .disconnected { color: #ff4d4f; font-size: 12px; margin-left: 10px; }
    </style>
</head>
<body>
    <div class="container">
        <h1>OpenCode Monitor <span id="connectionStatus" class="disconnected">● 未连接</span></h1>
        <div class="status-grid" id="statusGrid">
            <div class="status-card">
                <h2>当前状态</h2>
                <div class="status-value" id="currentStatus">等待中...</div>
                <div class="status-time" id="statusTime"></div>
            </div>
        </div>
        <div class="log">
            <h2>状态日志</h2>
            <div class="log-list" id="logList"></div>
        </div>
    </div>

    <script>
        const statusGrid = document.getElementById('statusGrid');
        const logList = document.getElementById('logList');
        const currentStatus = document.getElementById('currentStatus');
        const statusTime = document.getElementById('statusTime');
        const connectionStatus = document.getElementById('connectionStatus');
        
        let ws = null;
        let reconnectTimer = null;

        function connect() {
            const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
            ws = new WebSocket(protocol + '//' + window.location.host + '/ws');

            ws.onopen = function() {
                connectionStatus.textContent = '● 已连接';
                connectionStatus.className = 'connected';
                if (reconnectTimer) {
                    clearTimeout(reconnectTimer);
                    reconnectTimer = null;
                }
            };

            ws.onmessage = function(event) {
                try {
                    const data = JSON.parse(event.data);
                    updateStatus(data);
                    addLog(data);
                } catch (e) {
                    console.error('解析消息失败:', e);
                }
            };

            ws.onclose = function() {
                connectionStatus.textContent = '● 未连接';
                connectionStatus.className = 'disconnected';
                reconnectTimer = setTimeout(connect, 3000);
            };

            ws.onerror = function() {
                ws.close();
            };
        }

        function updateStatus(data) {
            currentStatus.textContent = data.status;
            currentStatus.className = 'status-value status-' + data.status;
            statusTime.textContent = '端口: ' + data.port + ' | 更新时间: ' + new Date().toLocaleTimeString();
        }

        function addLog(data) {
            const item = document.createElement('div');
            item.className = 'log-item';
            item.innerHTML = '<span class="log-time">' + new Date().toLocaleTimeString() + '</span>' +
                           '<span class="log-port">[:' + data.port + ']</span>' +
                           '<span class="status-' + data.status + '">' + data.status + '</span>';
            logList.insertBefore(item, logList.firstChild);
            
            while (logList.children.length > 50) {
                logList.removeChild(logList.lastChild);
            }
        }

        connect();
    </script>
</body>
</html>`

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Write([]byte(html))
}