Cyrene/backend/gateway/internal/handler/chat_handler.go

package handler

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"github.com/yourname/cyrene-ai/pkg/logger"
	"net/http"
	"strings"
	"time"

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

	"github.com/yourname/cyrene-ai/gateway/internal/config"
	"github.com/yourname/cyrene-ai/gateway/internal/store"
	"github.com/yourname/cyrene-ai/gateway/internal/ws"
)

// ChatHandler 聊天处理器
type ChatHandler struct {
	cfg          *config.Config
	hub          *ws.Hub
	sessionStore *store.SessionStore
	upgrader     websocket.Upgrader
}

// NewChatHandler 创建聊天处理器
func NewChatHandler(cfg *config.Config, hub *ws.Hub, sessionStore *store.SessionStore) *ChatHandler {
	return &ChatHandler{
		cfg:          cfg,
		hub:          hub,
		sessionStore: sessionStore,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin: func(r *http.Request) bool {
				return true // 开发阶段允许所有来源
			},
		},
	}
}

// HandleWebSocket 处理WebSocket升级和消息路由
func (h *ChatHandler) HandleWebSocket(c *gin.Context) {
	// 从query参数获取token和session_id
	token := c.Query("token")
	sessionID := c.Query("session_id")

	if token == "" {
		// 也尝试从Authorization头读取
		authHeader := c.GetHeader("Authorization")
		if len(authHeader) > 7 && authHeader[:7] == "Bearer " {
			token = authHeader[7:]
		}
	}

	if token == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "需要认证令牌"})
		return
	}

	// 验证token
	userID, err := h.cfg.ValidateToken(token)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "认证令牌无效"})
		return
	}

	// 主对话仅限管理员访问
	if userID != "admin" {
		c.JSON(http.StatusForbidden, gin.H{
			"error":     "主对话仅限管理员使用",
			"errorType": "admin_only",
			"hint":      "请使用管理员账号登录以访问主对话功能",
		})
		return
	}

	if sessionID == "" {
		sessionID = "session_" + generateID()
	}

	// 升级WebSocket连接
	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.Printf("[WS] 升级连接失败: %v", err)
		return
	}

	// 创建客户端
	client := ws.NewClient(h.hub, conn, userID, sessionID)

	// 注册到Hub
	h.hub.Register(client)

	// 启动读写协程
	go client.WritePump()
	go client.ReadPump(func(client *ws.Client, msg ws.ClientMessage) {
		h.handleMessage(client, msg)
	})
}

// handleMessage 处理WebSocket消息
func (h *ChatHandler) handleMessage(client *ws.Client, msg ws.ClientMessage) {
	switch msg.Type {
	case "message":
		h.handleChatMessage(client, msg)
	case "voice_input":
		h.handleVoiceInput(client, msg)
	case "history":
		h.handleHistoryRequest(client, msg)
	default:
		logger.Printf("[WS] 未知消息类型: %s from user=%s", msg.Type, client.UserID)
	}
}

// handleChatMessage 处理文字聊天消息 - 转发到 AI-Core（流式发送）
func (h *ChatHandler) handleChatMessage(client *ws.Client, msg ws.ClientMessage) {
	mode := msg.Mode
	if mode == "" {
		mode = "text"
	}

	// 持久化用户消息到数据库（在 WebSocket 发送之前）
	if h.sessionStore != nil && h.sessionStore.IsAvailable() {
		if err := h.sessionStore.AddMessage(client.SessionID, "user", "chat", msg.Content); err != nil {
			logger.Printf("[chat] 持久化用户消息失败: %v", err)
		}
	}

	// 记录用户消息
	h.hub.RecordMessage(client.SessionID, "user", msg.Content)

	// 设置会话状态为 thinking
	h.hub.UpdateSessionState(client.SessionID, "thinking")

	// 构建 AI-Core 请求
	aiReq := map[string]interface{}{
		"user_id":    client.UserID,
		"session_id": client.SessionID,
		"message":    msg.Content,
		"mode":       mode,
	}
	if len(msg.Attachments) > 0 {
		aiReq["attachments"] = msg.Attachments
	}
	reqBody, err := json.Marshal(aiReq)
	if err != nil {
		logger.Printf("[chat] 序列化请求失败: %v", err)
		h.hub.UpdateSessionState(client.SessionID, "error")
		client.SendMessage(ws.ServerMessage{
			Type:      "error",
			MessageID: "msg_" + generateID(),
			Error:     "内部错误，请稍后重试",
			Timestamp: time.Now().UnixMilli(),
		})
		return
	}

	// 缓存用户消息（在 goroutine 前完成，避免竞态）
	userMsg := ws.Message{
		ID:        "msg_" + generateID(),
		Role:      "user",
		Content:   msg.Content,
		Timestamp: time.Now().UnixMilli(),
	}
	if len(msg.Attachments) > 0 {
		userMsg.Attachments = msg.Attachments
	}
	h.hub.CacheMessage(client.UserID, client.SessionID, userMsg)

	// 在 goroutine 中进行 AI-Core 调用和流式发送，避免阻塞 ReadPump
	go h.streamResponse(client, mode, reqBody, msg.Content)
}

// streamResponse 调用 AI-Core SSE 流式接口并逐 delta 转发给客户端
func (h *ChatHandler) streamResponse(client *ws.Client, mode string, reqBody []byte, userMsg string) {
	aiCoreURL := h.cfg.AICoreURL + "/api/v1/chat"
	httpReq, err := http.NewRequest("POST", aiCoreURL, bytes.NewReader(reqBody))
	if err != nil {
		logger.Printf("[chat] 创建 AI-Core 请求失败: %v", err)
		h.hub.UpdateSessionState(client.SessionID, "error")
		client.SendMessage(ws.ServerMessage{
			Type:      "error",
			MessageID: "msg_" + generateID(),
			Error:     "服务暂不可用",
			Timestamp: time.Now().UnixMilli(),
		})
		return
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "text/event-stream")

	httpClient := &http.Client{Timeout: 120 * time.Second}
	resp, err := httpClient.Do(httpReq)
	if err != nil {
		logger.Printf("[chat] AI-Core 调用失败: %v", err)
		h.hub.UpdateSessionState(client.SessionID, "error")
		client.SendMessage(ws.ServerMessage{
			Type:      "error",
			MessageID: "msg_" + generateID(),
			Error:     fmt.Sprintf("AI-Core 调用失败: %v", err),
			Timestamp: time.Now().UnixMilli(),
		})
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		logger.Printf("[chat] AI-Core 返回错误 [%d]: %s", resp.StatusCode, string(body))
		h.hub.UpdateSessionState(client.SessionID, "error")
		client.SendMessage(ws.ServerMessage{
			Type:      "error",
			MessageID: "msg_" + generateID(),
			Error:     fmt.Sprintf("AI-Core 错误 (%d)", resp.StatusCode),
			Timestamp: time.Now().UnixMilli(),
		})
		return
	}

	// 使用 bufio.Scanner 逐行读取 SSE 响应
	scanner := bufio.NewScanner(resp.Body)
	// 增大 scanner buffer 以处理大块 SSE 数据
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

		// 通知前端 AI 开始生成回复
		client.SendMessage(ws.ServerMessage{
			Type:      "stream_start",
			MessageID: "msg_" + generateID(),
			SessionID: client.SessionID,
			Timestamp: time.Now().UnixMilli(),
		})

	var fullText string
	var msgID string
	var hasReview bool             // 是否有审查消息（避免重复持久化）
	var segments []ws.VoiceSegment // 收集断句信息

	for scanner.Scan() {
		line := scanner.Text()

		// 跳过非 data 行
		if !strings.HasPrefix(line, "data: ") {
			continue
		}

		data := strings.TrimPrefix(line, "data: ")

		// SSE 流结束标记
		if data == "[DONE]" {
			break
		}

		// 解析 delta 数据
		var chunk struct {
			Delta     string `json:"delta"`
			Error     string `json:"error,omitempty"`
			MessageID string `json:"message_id,omitempty"`
			Mode      string `json:"mode,omitempty"`
			Done      bool   `json:"done,omitempty"`
			// 断句相关 (来自 AI-Core 新格式)
			Segments []struct {
				Index int    `json:"index"`
				Text  string `json:"text"`
			} `json:"segments,omitempty"`
			// 审查后的结构化消息
			ReviewMessages []ws.ReviewMessage `json:"review_messages,omitempty"`
		}
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			logger.Printf("[chat] 解析 SSE delta 失败: %v, raw=%s", err, data)
			continue
		}

		// 错误处理
		if chunk.Error != "" {
			logger.Printf("[chat] AI-Core 流式错误: %s", chunk.Error)
			h.hub.UpdateSessionState(client.SessionID, "error")
			client.SendMessage(ws.ServerMessage{
				Type:      "error",
				MessageID: "msg_" + generateID(),
				Error:     chunk.Error,
				Timestamp: time.Now().UnixMilli(),
			})
			return
		}

		// 记录 message_id
		if chunk.MessageID != "" {
			msgID = chunk.MessageID
		}

		// 如果是结束标记（含 done: true），跳出
		if chunk.Done {
			break
		}

		// 处理审查后的结构化消息 (review)
		if len(chunk.ReviewMessages) > 0 {
			for i, rm := range chunk.ReviewMessages {
				role := "assistant"
				msgType := "chat"
				if rm.Type == "action" {
					role = "action"
					msgType = "action"
				}
				reviewMsgID := fmt.Sprintf("%s_r%d", msgID, i)
				// 持久化每条审查消息
				if h.sessionStore != nil && h.sessionStore.IsAvailable() {
					if err := h.sessionStore.AddMessage(client.SessionID, role, msgType, rm.Content); err != nil {
						logger.Printf("[chat] 持久化审查消息失败: %v", err)
					}
				}
				h.hub.CacheMessage(client.UserID, client.SessionID, ws.Message{
					ID:        reviewMsgID,
					Role:      role,
					Content:   rm.Content,
					Timestamp: time.Now().UnixMilli(),
				})
				client.SendMessage(ws.ServerMessage{
					Type:           "response",
					MessageID:      reviewMsgID,
					Content:        rm.Content,
					Role:           role,
					MsgType:        msgType,
					SessionID:      client.SessionID,
					Timestamp:      time.Now().UnixMilli(),
				})
				// 使用 MessageScheduler 计算的 per-message 延迟
				if rm.DelayMs > 0 {
					time.Sleep(time.Duration(rm.DelayMs) * time.Millisecond)
				}
			}
			hasReview = true
			continue
		}

		// 处理断句事件 (stream_segments)
		if len(chunk.Segments) > 0 {
			for _, seg := range chunk.Segments {
				segments = append(segments, ws.VoiceSegment{
					Index: seg.Index,
					Text:  seg.Text,
				})
			}
			// 发送断句事件给前端
			client.SendMessage(ws.ServerMessage{
				Type:      "stream_segments",
				MessageID: msgID,
				Segments:  segments,
				SessionID: client.SessionID,
				Timestamp: time.Now().UnixMilli(),
			})
			continue
		}

		// 逐 delta 积累（不再逐块转发，由审查消息代替）
		if chunk.Delta != "" {
			fullText += chunk.Delta
		}
	}

	if err := scanner.Err(); err != nil {
		logger.Printf("[chat] SSE 读取错误: %v", err)
		h.hub.UpdateSessionState(client.SessionID, "error")
		client.SendMessage(ws.ServerMessage{
			Type:      "error",
			MessageID: "msg_" + generateID(),
			Error:     fmt.Sprintf("流读取错误: %v", err),
			Timestamp: time.Now().UnixMilli(),
		})
		return
	}

	if msgID == "" {
		msgID = "msg_" + generateID()
	}

	// 检测是否为多消息格式（包含空行分隔的多条消息）
	// 如果已有审查消息则跳过，避免与 review_messages 重复
	multiParts := parseMultiMessage(fullText)
	if !hasReview && len(multiParts) > 1 {
		// 发送 multi_message 事件
		var items []ws.MultiMessageItem
		for i, part := range multiParts {
			items = append(items, ws.MultiMessageItem{
				Index:   i,
				Content: part,
			})
		}
		client.SendMessage(ws.ServerMessage{
			Type:      "multi_message",
			MessageID: msgID,
			SessionID: client.SessionID,
			MultiMessage: &ws.MultiMessagePayload{
				Messages: items,
			},
			Timestamp: time.Now().UnixMilli(),
		})
	}

	// 发送 stream_end
	client.SendMessage(ws.ServerMessage{
		Type:      "stream_end",
		MessageID: msgID,
		SessionID: client.SessionID,
		Timestamp: time.Now().UnixMilli(),
	})

	// 持久化 AI 回复到数据库
	// 如果有审查消息，每条已单独持久化，跳过 fullText 以避免重复
	if !hasReview && fullText != "" {
		if h.sessionStore != nil && h.sessionStore.IsAvailable() {
			if err := h.sessionStore.AddMessage(client.SessionID, "assistant", "chat", fullText); err != nil {
				logger.Printf("[chat] 持久化 AI 回复失败: %v", err)
			}
		}

		h.hub.CacheMessage(client.UserID, client.SessionID, ws.Message{
			ID:        msgID,
			Role:      "assistant",
			Content:   fullText,
			Timestamp: time.Now().UnixMilli(),
		})
	}
	// RecordMessage 使用不带 [review] 标记的文本
	recordText := fullText
	if hasReview {
		recordText = strings.ReplaceAll(fullText, "[review]", "")
	}
	h.hub.RecordMessage(client.SessionID, "assistant", recordText)

	// 设置会话状态为 idle
	h.hub.UpdateSessionState(client.SessionID, "idle")
}

// handleVoiceInput 处理语音输入
func (h *ChatHandler) handleVoiceInput(client *ws.Client, msg ws.ClientMessage) {
	// MVP阶段：返回提示
	response := ws.ServerMessage{
		Type:      "error",
		MessageID: "msg_" + generateID(),
		Error:     "语音处理功能将在后续版本中启用",
		Timestamp: time.Now().UnixMilli(),
	}
	client.SendMessage(response)
}


// handleHistoryRequest 处理历史消息请求
func (h *ChatHandler) handleHistoryRequest(client *ws.Client, msg ws.ClientMessage) {
	// 优先使用请求中的 session_id，否则使用客户端的 session_id
	sessionID := msg.SessionID
	if sessionID == "" {
		sessionID = client.SessionID
	}

	messages := h.hub.GetConversation(client.UserID, sessionID)

	// 如果内存缓存为空，尝试从数据库恢复（网关重启后缓存丢失的情况）
	if len(messages) == 0 && h.sessionStore != nil && h.sessionStore.IsAvailable() {
		dbMessages, err := h.sessionStore.GetMessages(sessionID, 50, 0)
		if err == nil && len(dbMessages) > 0 {
			logger.Printf("[history] 从数据库恢复会话历史: session=%s, %d 条消息", sessionID, len(dbMessages))
			// 恢复到内存缓存
			for _, dbMsg := range dbMessages {
				messages = append(messages, ws.Message{
					ID:        fmt.Sprintf("db_%d", dbMsg.ID),
					Role:      dbMsg.Role,
					MsgType:   dbMsg.MsgType,
					Content:   dbMsg.Content,
					Timestamp: dbMsg.CreatedAt.UnixMilli(),
				})
				h.hub.CacheMessage(client.UserID, sessionID, ws.Message{
					ID:        fmt.Sprintf("db_%d", dbMsg.ID),
					Role:      dbMsg.Role,
					MsgType:   dbMsg.MsgType,
					Content:   dbMsg.Content,
					Timestamp: dbMsg.CreatedAt.UnixMilli(),
				})
			}
		}
	}

	if messages == nil {
		messages = []ws.Message{}
	}

	response := ws.ServerMessage{
		Type:      "history_response",
		MessageID: "hist_" + generateID(),
		Messages:  messages,
		Timestamp: time.Now().UnixMilli(),
	}

	if err := client.SendMessage(response); err != nil {
		logger.Printf("[WS] 发送历史消息失败: %v", err)
	}
}

// SendSystemMessage 向用户发送系统消息（用于主动通知）
func (h *ChatHandler) SendSystemMessage(userID, sessionID, text string) error {
	msg := ws.ServerMessage{
		Type:      "response",
		MessageID: "sys_" + generateID(),
		Text:      text,
		Timestamp: time.Now().UnixMilli(),
	}

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

	h.hub.SendToSession(userID, sessionID, data)
	return nil
}

// HandleProactiveMessage 处理来自 AI-Core 后台思考的主动消息
// POST /api/v1/internal/proactive-message
func (h *ChatHandler) HandleProactiveMessage(c *gin.Context) {
	var req struct {
		UserID    string `json:"user_id" binding:"required"`
		Content   string `json:"content" binding:"required"`
		SessionID string `json:"session_id"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数无效: " + err.Error()})
		return
	}

	// 检查用户是否在线
	onlineCount := h.hub.UserClientCount(req.UserID)
	if onlineCount == 0 {
		// Phase 2: 离线时排队，等待用户重连后推送
		data, _ := json.Marshal(ws.ServerMessage{
			Type:      "response",
			MessageID: "proactive_" + generateID(),
			Content:   req.Content,
			Role:      "assistant",
			MsgType:   "proactive",
			SessionID: req.SessionID,
			Timestamp: time.Now().UnixMilli(),
		})
		h.hub.QueueProactiveMessage(req.UserID, data)
		logger.Printf("[proactive] 用户离线，消息已排队: user=%s", req.UserID)
		c.JSON(http.StatusOK, gin.H{
			"success":  true,
			"reason":   "queued",
			"message":  "用户离线，消息已排队等待重连后推送",
		})
		return
	}

	// 构建主动消息
	msgID := "proactive_" + generateID()
	msg := ws.ServerMessage{
		Type:      "response",
		MessageID: msgID,
		Content:   req.Content,
		Role:      "assistant",
		MsgType:   "proactive",
		SessionID: req.SessionID,
		Timestamp: time.Now().UnixMilli(),
	}

	data, err := json.Marshal(msg)
	if err != nil {
		logger.Printf("[proactive] 序列化消息失败: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "内部错误"})
		return
	}

	h.hub.SendToUser(req.UserID, data)

	// 同时缓存到对话历史（使用 admin 的主 session）
	sessionID := req.SessionID
	if sessionID == "" {
		sessionID = "session_admin_main"
	}
	h.hub.CacheMessage(req.UserID, sessionID, ws.Message{
		ID:        msgID,
		Role:      "assistant",
		Content:   req.Content,
		Timestamp: time.Now().UnixMilli(),
	})
	h.hub.RecordMessage(sessionID, "assistant", req.Content)

	logger.Printf("[proactive] 主动消息已推送: user=%s, online=%d, content_len=%d", req.UserID, onlineCount, len(req.Content))

	c.JSON(http.StatusOK, gin.H{
		"success":   true,
		"message":   "消息已推送",
		"delivered": onlineCount,
	})
}

func generateID() string {
	return time.Now().Format("20060102150405") + randomStr(6)
}

func randomStr(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		// fallback: deterministic but hard to predict
		for i := range b {
			b[i] = byte(time.Now().UnixNano()%256)
		}
	}
	return hex.EncodeToString(b)[:n]
}

// parseMultiMessage 检测并解析多消息格式
// 如果文本包含空行分隔的多条消息，拆分为多条；否则返回单条
func parseMultiMessage(text string) []string {
	if text == "" {
		return nil
	}
	// 按双换行（空行）分割
	parts := strings.Split(text, "\n\n")
	// 过滤空字符串并去除首尾空白
	var result []string
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p != "" {
			result = append(result, p)
		}
	}
	// 如果只有一条，返回 nil 表示不是多消息格式
	if len(result) <= 1 {
		return nil
	}
	return result
}