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

package api

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

	"github.com/gorilla/websocket"

	"ai-status-light/internal/database"
	"ai-status-light/internal/logger"
	"ai-status-light/internal/web"
)

type ClientStatus struct {
	Port      int    `json:"port"`
	Status    string `json:"status"`
	Code      string `json:"code"`
	Timestamp string `json:"timestamp"`
}

type Server struct {
	db        *database.DB
	server    *http.Server
	clients   map[*websocket.Conn]bool
	clientsMu sync.Mutex
	upgrader  websocket.Upgrader
	statusMap map[int]*ClientStatus
	statusMu  sync.RWMutex
	certFile  string
	keyFile   string
}

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),
		statusMap: make(map[int]*ClientStatus),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/api/clients", s.handleClients)
	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("/", web.Handler())

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

func (s *Server) EnableTLS(certFile, keyFile string) {
	s.certFile = certFile
	s.keyFile = keyFile
}

func (s *Server) Start() error {
	if s.certFile != "" && s.keyFile != "" {
		logger.Info("API 服务器开始监听 (HTTPS): %s", s.server.Addr)
		err := s.server.ListenAndServeTLS(s.certFile, s.keyFile)
		if err != nil && err != http.ErrServerClosed {
			logger.Error("API 服务器监听失败: %v", err)
		}
		return err
	}
	logger.Info("API 服务器开始监听: %s", s.server.Addr)
	err := s.server.ListenAndServe()
	if err != nil && err != http.ErrServerClosed {
		logger.Error("API 服务器监听失败: %v", err)
	}
	return err
}

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) handleClients(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeJSON(w, http.StatusMethodNotAllowed, Response{Code: -1, Message: "方法不允许"})
		return
	}

	s.statusMu.RLock()
	result := make([]*ClientStatus, 0, len(s.statusMap))
	for _, cs := range s.statusMap {
		result = append(result, cs)
	}
	s.statusMu.RUnlock()

	sort.Slice(result, func(i, j int) bool {
		return result[i].Port < result[j].Port
	})

	writeJSON(w, http.StatusOK, Response{Code: 0, Message: "success", Data: result})
}

func (s *Server) handleMQTT(w http.ResponseWriter, r *http.Request) {
	logger.Debug("HTTP %s %s", r.Method, r.URL.Path)
	switch r.Method {
	case http.MethodGet:
		s.listMQTTConfigs(w, r)
	case http.MethodPost:
		s.createMQTTConfig(w, r)
	default:
		logger.Warn("不支持的 HTTP 方法: %s %s", r.Method, r.URL.Path)
		writeJSON(w, http.StatusMethodNotAllowed, Response{Code: -1, Message: "方法不允许"})
	}
}

func (s *Server) handleMQTTByID(w http.ResponseWriter, r *http.Request) {
	logger.Debug("HTTP %s %s", r.Method, r.URL.Path)
	idStr := strings.TrimPrefix(r.URL.Path, "/api/mqtt/")
	id, err := strconv.Atoi(idStr)
	if err != nil {
		logger.Warn("无效的配置 ID: %s", idStr)
		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:
		logger.Warn("不支持的 HTTP 方法: %s %s", r.Method, r.URL.Path)
		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 {
		logger.Error("查询 MQTT 配置列表失败: %v", err)
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}
	logger.Debug("查询 MQTT 配置列表: %d 条", len(configs))
	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 {
		logger.Warn("创建配置请求体解析失败: %v", err)
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "无效的请求体"})
		return
	}

	if cfg.Broker == "" {
		logger.Warn("创建配置: 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 {
		logger.Error("创建 MQTT 配置失败: %v", err)
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	logger.Info("MQTT 配置已创建: id=%d, broker=%s, topic=%s", cfg.ID, cfg.Broker, cfg.Topic)
	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 {
		logger.Error("查询 MQTT 配置失败: id=%d, %v", id, err)
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	for _, cfg := range configs {
		if cfg.ID == id {
			logger.Debug("查询 MQTT 配置: id=%d", id)
			writeJSON(w, http.StatusOK, Response{Code: 0, Message: "ok", Data: cfg})
			return
		}
	}

	logger.Warn("MQTT 配置不存在: id=%d", id)
	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 {
		logger.Warn("更新配置请求体解析失败: id=%d, %v", id, err)
		writeJSON(w, http.StatusBadRequest, Response{Code: -1, Message: "无效的请求体"})
		return
	}

	cfg.ID = id
	if cfg.Broker == "" {
		logger.Warn("更新配置: broker 为空, id=%d", id)
		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 {
		logger.Error("更新 MQTT 配置失败: id=%d, %v", id, err)
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}

	logger.Info("MQTT 配置已更新: id=%d, broker=%s, topic=%s", id, cfg.Broker, cfg.Topic)
	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 {
		logger.Error("删除 MQTT 配置失败: id=%d, %v", id, err)
		writeJSON(w, http.StatusInternalServerError, Response{Code: -1, Message: err.Error()})
		return
	}
	logger.Info("MQTT 配置已删除: id=%d", id)
	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 {
		logger.Error("WebSocket 升级失败: %v", err)
		return
	}

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

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

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

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

func (s *Server) BroadcastStatus(port int, status string, code string) {
	ts := time.Now().Format(time.RFC3339)

	s.statusMu.Lock()
	s.statusMap[port] = &ClientStatus{
		Port:      port,
		Status:    status,
		Code:      code,
		Timestamp: ts,
	}
	s.statusMu.Unlock()

	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()

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

	payload := map[string]interface{}{
		"port":      port,
		"status":    status,
		"code":      code,
		"timestamp": ts,
	}

	data, err := json.Marshal(payload)
	if err != nil {
		logger.Error("序列化广播消息失败: %v", err)
		return
	}

	for client := range s.clients {
		err := client.WriteMessage(websocket.TextMessage, data)
		if err != nil {
			logger.Debug("WebSocket 写入失败，移除客户端: %v", err)
			client.Close()
			delete(s.clients, client)
		}
	}
}