CamTalk/backend/internal/ws/handler.go

package ws

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"

	"github.com/hhs/camtalk/internal/logger"
	"github.com/hhs/camtalk/internal/models"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true }, // 开发阶段允许所有来源
}

// Client 代表一个 WebSocket 客户端连接。
type Client struct {
	conn      *websocket.Conn
	sessionID string
	mu        sync.Mutex
}

func (c *Client) sendJSON(v any) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.conn.WriteJSON(v)
}

// ServeWS 处理 WebSocket 升级请求。
func ServeWS(c *gin.Context) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.Log.Errorw("websocket upgrade failed", "error", err)
		return
	}
	defer conn.Close()

	sessionID := uuid.New().String()
	client := &Client{conn: conn, sessionID: sessionID}

	// 发送 connected 消息
	_ = client.sendJSON(models.WsConnected{
		Type:          "connected",
		SessionID:     sessionID,
		ServerVersion: "0.1.0",
	})
	logger.Log.Infow("client connected", "session", sessionID)

	// 心跳检测
	lastPong := time.Now()
	conn.SetPongHandler(func(string) error {
		lastPong = time.Now()
		return nil
	})

	// 启动心跳检查 goroutine
	done := make(chan struct{})
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if time.Since(lastPong) > 60*time.Second {
					logger.Log.Warnw("heartbeat timeout", "session", sessionID)
					conn.Close()
					return
				}
			case <-done:
				return
			}
		}
	}()

	// 消息读取循环
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				logger.Log.Warnw("ws read error", "error", err)
			}
			break
		}

		// 解析消息类型
		var envelope struct {
			Type string `json:"type"`
		}
		if err := json.Unmarshal(message, &envelope); err != nil {
			_ = client.sendJSON(models.WsError{
				Type:    "error",
				Code:    "INVALID_MESSAGE",
				Message: "invalid JSON",
			})
			continue
		}

		switch envelope.Type {
		case "ping":
			_ = client.sendJSON(models.WsPong{Type: "pong"})

		case "query":
			var msg models.WsQuery
			if err := json.Unmarshal(message, &msg); err != nil {
				_ = client.sendJSON(models.WsError{
					Type:      "error",
					Code:      "INVALID_MESSAGE",
					Message:   "invalid query message",
					RequestID: msg.RequestID,
				})
				continue
			}
			logger.Log.Infow("query received", "session", sessionID, "request", msg.RequestID)
			// TODO: 调用 AI 编排流程（STT → LLM → TTS）

		case "config":
			var msg models.WsConfig
			if err := json.Unmarshal(message, &msg); err != nil {
				_ = client.sendJSON(models.WsError{
					Type:    "error",
					Code:    "INVALID_MESSAGE",
					Message: "invalid config message",
				})
				continue
			}
			logger.Log.Infow("config update", "session", sessionID)
			// TODO: 更新会话配置

		case "interrupt":
			logger.Log.Infow("interrupt received", "session", sessionID)
			// TODO: 中断当前 AI 响应

		default:
			_ = client.sendJSON(models.WsError{
				Type:    "error",
				Code:    "INVALID_MESSAGE",
				Message: "unknown message type: " + envelope.Type,
			})
		}
	}

	close(done)
	logger.Log.Infow("client disconnected", "session", sessionID)
}