CamTalk/backend/internal/ws/handler.go

package ws

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

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

	"github.com/hhs/camtalk/internal/ai/llm"
	"github.com/hhs/camtalk/internal/auth"
	"github.com/hhs/camtalk/internal/config"
	"github.com/hhs/camtalk/internal/errors"
	"github.com/hhs/camtalk/internal/models"
	"github.com/hhs/camtalk/internal/orchestrator"
	"github.com/hhs/camtalk/internal/ratelimit"
	"github.com/hhs/camtalk/internal/session"
	"github.com/hhs/camtalk/internal/store"
	"github.com/hhs/camtalk/internal/trace"
)

// newUpgrader 根据配置创建 WebSocket upgrader。
func newUpgrader(cfg *config.Config) websocket.Upgrader {
	allowedOrigins := cfg.Server.AllowedOrigins
	return websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			if len(allowedOrigins) == 0 {
				return true // 未配置则允许所有来源（开发模式）
			}
			origin := r.Header.Get("Origin")
			for _, o := range allowedOrigins {
				if o == origin || o == "*" {
					return true
				}
			}
			return false
		},
	}
}

// Client 代表一个 WebSocket 客户端连接。
type Client struct {
	conn         *websocket.Conn
	sessionID    string
	sessionMgr   session.Manager
	orchestrator orchestrator.Orchestrator
	cancelFuncs  map[string]context.CancelFunc // requestID → cancel func
	mu           sync.Mutex
}

// SendJSON 向客户端发送 JSON 消息（公开以便 errors 包调用）。
func (c *Client) SendJSON(v any) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.conn.WriteJSON(v)
}

// WSClient 实现 orchestrator.Sender 接口，将消息推送到 WebSocket 连接。
type WSClient struct {
	client    *Client
	requestID string
}

// SendSTTResult 发送语音识别结果。
func (w *WSClient) SendSTTResult(result models.WsSTTResult) error {
	result.RequestID = w.requestID
	return w.client.SendJSON(result)
}

// SendLLMChunk 发送 LLM 流式文本增量。
func (w *WSClient) SendLLMChunk(chunk models.WsLLMChunk) error {
	chunk.RequestID = w.requestID
	return w.client.SendJSON(chunk)
}

// SendLLMDone 发送 LLM 流结束信号。
func (w *WSClient) SendLLMDone(done models.WsLLMDone) error {
	done.RequestID = w.requestID
	return w.client.SendJSON(done)
}

// SendTTSAudio 发送 TTS 音频数据。
func (w *WSClient) SendTTSAudio(audio models.WsTTSAudio) error {
	audio.RequestID = w.requestID
	return w.client.SendJSON(audio)
}

// SendError 发送错误消息。
func (w *WSClient) SendError(err models.WsError) error {
	err.RequestID = w.requestID
	return w.client.SendJSON(err)
}

// ServeWS 处理 WebSocket 升级请求。
func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *config.Config, tokenMgr *auth.TokenManager, limiter ratelimit.Limiter, scenarioRepo store.UserScenarioRepository) gin.HandlerFunc {
	upgrader := newUpgrader(cfg)
	heartbeatInterval := time.Duration(cfg.Server.HeartbeatInterval) * time.Second
	heartbeatTimeout := time.Duration(cfg.Server.HeartbeatTimeout) * time.Second
	version := cfg.App.Version

	return func(c *gin.Context) {
		serveWS(c, sessionMgr, orch, upgrader, heartbeatInterval, heartbeatTimeout, version, tokenMgr, limiter, scenarioRepo)
	}
}

func serveWS(c *gin.Context, sessionMgr session.Manager, orch orchestrator.Orchestrator,
	upgrader websocket.Upgrader, heartbeatInterval, heartbeatTimeout time.Duration, version string, tokenMgr *auth.TokenManager, limiter ratelimit.Limiter, scenarioRepo store.UserScenarioRepository) {

	// --- JWT 认证（upgrade 前完成，失败直接返回 HTTP 错误） ---
	token := c.Query("token")
	if token == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing token"})
		return
	}
	claims, err := tokenMgr.ValidateAccess(token)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
		return
	}
	userID := claims.UserID
	username := claims.Username

	// --- conversation_id 处理（upgrade 前校验归属） ---
	conversationID := c.Query("conversation_id")
	if conversationID != "" {
		sess, err := sessionMgr.Get(c.Request.Context(), conversationID)
		if err != nil || sess.UserID != userID {
			c.JSON(http.StatusUnauthorized, gin.H{"error": "SESSION_NOT_FOUND"})
			return
		}
	}

	// 生成连接级 trace ID（整个 WebSocket 生命周期使用）
	ctx := c.Request.Context()
	traceID := trace.GetTraceID(ctx)
	if traceID == "" {
		// 如果 REST 中间件未生成（不应发生），fallback 生成
		traceID = trace.GenerateTraceID()
		ctx = trace.WithTraceID(ctx, traceID)
		c.Request = c.Request.WithContext(ctx)
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log := trace.FromContext(ctx)
		log.Errorw("websocket upgrade failed", "error", err)
		return
	}
	defer conn.Close()

	// 创建或复用会话
	var sessionID string
	if conversationID != "" {
		sessionID = conversationID
		ctx = trace.WithSessionID(ctx, sessionID)
		log := trace.FromContext(ctx)
		log.Infow("resuming conversation", "user_id", userID)
	} else {
		sessionID, err = sessionMgr.Create(context.Background(), userID, models.DefaultConfig())
		if err != nil {
			log := trace.FromContext(ctx)
			log.Errorw("create session failed", "error", err)
			return
		}
		ctx = trace.WithSessionID(ctx, sessionID)
	}

	client := &Client{
		conn:         conn,
		sessionID:    sessionID,
		sessionMgr:   sessionMgr,
		orchestrator: orch,
		cancelFuncs:  make(map[string]context.CancelFunc),
	}

	// 发送 connected 消息
	_ = client.SendJSON(models.WsConnected{
		Type:          "connected",
		SessionID:     sessionID,
		ServerVersion: version,
	})
	log := trace.FromContext(ctx)
	log.Infow("client connected", "user_id", userID, "username", username)

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

	// 启动心跳检查 goroutine
	done := make(chan struct{})
	go func() {
		ticker := time.NewTicker(heartbeatInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if time.Since(lastPong) > heartbeatTimeout {
					log := trace.FromContext(ctx)
					log.Warnw("heartbeat timeout")
					conn.Close()
					return
				}
			case <-done:
				return
			}
		}
	}()

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

		// 解析消息类型
		var envelope struct {
			Type string `json:"type"`
		}
		if err := json.Unmarshal(message, &envelope); err != nil {
			errors.SendWSError(client, errors.CodeInvalidMessage, "", err)
			continue
		}

		switch envelope.Type {
		case "ping":
			lastPong = time.Now() // 刷新心跳计时器
			_ = client.SendJSON(models.WsPong{Type: "pong"})

		case "query":
			var msg models.WsQuery
			if err := json.Unmarshal(message, &msg); err != nil {
				errors.SendWSError(client, errors.CodeInvalidMessage, msg.RequestID, err)
				continue
			}

			// 注入 request ID 到 context
			queryCtx := trace.WithRequestID(ctx, msg.RequestID)
			log := trace.FromContext(queryCtx)
			log.Infow("query received", "has_image", msg.Image != "", "has_audio", msg.Audio != "")

			// 限流检查
			if limiter != nil {
				key := fmt.Sprintf("%s:query", userID)
				allowed, retryAfter := limiter.Allow(context.Background(), key)
				if !allowed {
					log.Warnw("rate limited", "user_id", userID, "retry_after", retryAfter)
					errors.SendWSError(client, errors.CodeRateLimited, msg.RequestID,
						fmt.Errorf("rate limited, retry after %s", retryAfter.Round(time.Second)))
					continue
				}
			}

			// 刷新会话 TTL
			if err := client.sessionMgr.Touch(context.Background(), sessionID); err != nil {
				log.Warnw("touch session failed", "error", err)
			}

			// 标记活跃请求
			if err := client.sessionMgr.SetActiveRequest(context.Background(), sessionID, msg.RequestID); err != nil {
				log.Warnw("set active request failed", "error", err)
			}

			// 创建可取消的 context
			processCtx, cancel := context.WithCancel(queryCtx)
			client.mu.Lock()
			client.cancelFuncs[msg.RequestID] = cancel
			client.mu.Unlock()

			// 创建 sender
			sender := &WSClient{client: client, requestID: msg.RequestID}

			// 启动 orchestrator 处理 goroutine
			go func() {
				defer func() {
					// 清理 cancel func
					client.mu.Lock()
					delete(client.cancelFuncs, msg.RequestID)
					client.mu.Unlock()
					cancel()
					// 清除活跃请求
					_ = client.sessionMgr.ClearActiveRequest(context.Background(), sessionID)
				}()

				if err := client.orchestrator.ProcessQuery(processCtx, sessionID, msg, sender); err != nil {
					log := trace.FromContext(processCtx)
					log.Errorw("process query failed", "error", err)
				}
			}()

		case "config":
			var msg models.WsConfig
			if err := json.Unmarshal(message, &msg); err != nil {
				errors.SendWSError(client, errors.CodeInvalidMessage, "", err)
				continue
			}

			patch := models.SessionConfigPatch{
				TTSEnabled:  msg.Payload.TTSEnabled,
				DetailLevel: msg.Payload.DetailLevel,
				Language:    msg.Payload.Language,
				Scenario:    msg.Payload.Scenario,
			}
			if err := client.sessionMgr.UpdateConfig(context.Background(), sessionID, patch); err != nil {
				errors.SendWSError(client, errors.CodeInternalError, "", err)
				continue
			}

			scenarioID := ""
			if msg.Payload.Scenario != nil {
				scenarioID = *msg.Payload.Scenario
			}
			log := trace.FromContext(ctx)
			log.Infow("config updated", "scenario", scenarioID)

			// 如果切换了情景（非自由对话），返回首句引导
			if scenarioID != "" && scenarioID != "free_chat" {
				sess, err := client.sessionMgr.Get(context.Background(), sessionID)
				if err == nil && sess != nil {
					// 加载用户自建情景
					var customGreetings map[string]string
					if sess.UserID != "" && scenarioRepo != nil {
						scenarios, err := scenarioRepo.FindByUserID(context.Background(), sess.UserID)
						if err == nil && len(scenarios) > 0 {
							customGreetings = make(map[string]string, len(scenarios))
							for _, s := range scenarios {
								if s.Greeting != "" {
									customGreetings[s.ID] = s.Greeting
								}
							}
						}
					}

					greeting := llm.GetScenarioGreeting(scenarioID, sess.Config.Language, customGreetings)
					if greeting != "" {
						// 发送首句作为 AI 消息
						_ = client.SendJSON(models.WsLLMChunk{
							Type:      "llm_chunk",
							RequestID: "scenario_greeting",
							Delta:     greeting,
							Role:      "assistant",
						})

						doneMsg := models.WsLLMDone{
							Type:      "llm_done",
							RequestID: "scenario_greeting",
							FullText:  greeting,
							Model:     "",
							LatencyMs: 0,
						}
						doneMsg.TokensUsed.Prompt = 0
						doneMsg.TokensUsed.Completion = 0
						doneMsg.TokensUsed.Total = 0
						_ = client.SendJSON(doneMsg)

						// 追加首句到历史记录
						_ = client.sessionMgr.AppendMessage(context.Background(), sessionID, models.Message{
							Role:    "assistant",
							Content: greeting,
						})
					}
				}
			}

		case "interrupt":
			log := trace.FromContext(ctx)
			log.Infow("interrupt received")

			// 获取活跃请求 ID 并取消
			reqID, _ := client.sessionMgr.GetActiveRequestID(context.Background(), sessionID)
			if reqID != "" {
				client.mu.Lock()
				if cancel, ok := client.cancelFuncs[reqID]; ok {
					cancel()
					delete(client.cancelFuncs, reqID)
				}
				client.mu.Unlock()
				_ = client.sessionMgr.ClearActiveRequest(context.Background(), sessionID)
			}

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

	close(done)

	// 取消所有活跃请求
	client.mu.Lock()
	for reqID, cancel := range client.cancelFuncs {
		log := trace.FromContext(ctx)
		log.Infow("canceling active request on disconnect", "request", reqID)
		cancel()
	}
	client.cancelFuncs = make(map[string]context.CancelFunc)
	client.mu.Unlock()

	// 断开连接时不销毁会话，让其自然过期（支持重连恢复）
	log = trace.FromContext(ctx)
	log.Infow("client disconnected")
}