gochat/internal/handler/api/v1/sse_stream_handler.go

package v1

import (
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/service"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// SSEStreamHandler handles Server-Sent Events streaming for Copilot real-time responses.
// Reference: M12 PRD §Copilot — Streaming SSE Support
//
// SSE event format:
//   event: message
//   data: {"content": "...", "done": false}
//
//   event: done
//   data: {"done": true}

// SSEStreamHandler holds dependencies for SSE streaming endpoints.
type SSEStreamHandler struct {
	copilotSvc  *service.CopilotService
	llmProvider llm.Provider
}

// NewSSEStreamHandler creates a new SSEStreamHandler.
func NewSSEStreamHandler(copilotSvc *service.CopilotService, llmProvider llm.Provider) *SSEStreamHandler {
	return &SSEStreamHandler{copilotSvc: copilotSvc, llmProvider: llmProvider}
}

// SSEStreamRequest holds the input for a streaming Copilot request.
type SSEStreamRequest struct {
	ThreadID uint   `json:"thread_id" validate:"required"`
	Content  string `json:"content" validate:"required,min=1"`
}

// StreamCopilotMessage streams a Copilot chat response via SSE.
// The OpenAIProvider.ChatCompletionStream uses a callback pattern (onChunk),
// so we write SSE events inside the callback as chunks arrive.
//
// POST /api/v1/accounts/:id/captain/copilot_threads/:thread_id/stream
func (h *SSEStreamHandler) StreamCopilotMessage(c *gin.Context) {
	accountID, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		writeSSEMessage(c, "error", `{"error": "invalid account_id"}`)
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}

	threadID, err := strconv.ParseUint(c.Param("thread_id"), 10, 64)
	if err != nil {
		writeSSEMessage(c, "error", `{"error": "invalid thread_id"}`)
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}

	var req SSEStreamRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeSSEMessage(c, "error", fmt.Sprintf(`{"error": "%s"}`, escapeJSONString(err.Error())))
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}

	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no") // disable nginx buffering

	// Get thread for context
	thread, err := h.copilotSvc.GetThreadByID(c.Request.Context(), uint(threadID))
	if err != nil {
		applogger.L().Errorf("SSE GetThread: %v", err)
		writeSSEMessage(c, "error", `{"error": "thread not found"}`)
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}

	if thread.AccountID != uint(accountID) {
		writeSSEMessage(c, "error", `{"error": "thread not found"}`)
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}
	if h.llmProvider == nil {
		writeSSEMessage(c, "error", `{"error": "Captain is disabled", "status": 422, "done": true}`)
		writeSSEMessage(c, "done", `{"done": true}`)
		return
	}

	// Build chat messages from thread history + new user message
	chatMessages := buildStreamChatMessages(thread, req.Content)

	// Stream from LLM using callback pattern
	streamReq := llm.ChatRequest{
		Model:       "gpt-4",
		Messages:    chatMessages,
		Temperature: 0.7,
		MaxTokens:   1024,
		Stream:      true,
	}

	err = h.llmProvider.ChatCompletionStream(c.Request.Context(), streamReq, func(chunk llm.StreamChunk) error {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			if content != "" {
				escaped := escapeJSONString(content)
				writeSSEMessage(c, "message", fmt.Sprintf(`{"content": "%s", "done": false}`, escaped))
				c.Writer.Flush()
			}

			finishReason := chunk.Choices[0].FinishReason
			if finishReason != "" && finishReason != "null" {
				writeSSEMessage(c, "done", `{"done": true, "content": ""}`)
				c.Writer.Flush()
				return io.EOF // signal end of stream
			}
		}
		return nil
	})

	if err != nil && err != io.EOF {
		applogger.L().Errorf("SSE stream error: %v", err)
		writeSSEMessage(c, "error", fmt.Sprintf(`{"error": "%s"}`, escapeJSONString(err.Error())))
		writeSSEMessage(c, "done", `{"done": true}`)
	}
}

// --- Helper Functions ---

// writeSSEMessage writes a single SSE event to the Gin response writer.
// SSE format: "event: <event>\ndata: <data>\n\n"
func writeSSEMessage(c *gin.Context, event string, data string) {
	c.Writer.WriteString(fmt.Sprintf("event: %s\ndata: %s\n\n", event, data))
}

// buildStreamChatMessages constructs chat messages for streaming from thread history + new content.
func buildStreamChatMessages(thread *model.CopilotThread, userContent string) []llm.ChatMessage {
	messages := []llm.ChatMessage{
		{Role: "system", Content: "You are an AI assistant helping a customer support agent. Provide helpful, professional, and concise responses."},
	}

	// Add thread history as context
	history := thread.PreviousHistory(thread.Messages)
	for _, h := range history {
		messages = append(messages, llm.ChatMessage{
			Role:    h.Role,
			Content: h.Content,
		})
	}

	// Add the new user message
	messages = append(messages, llm.ChatMessage{
		Role:    "user",
		Content: userContent,
	})

	return messages
}

// escapeJSONString escapes special characters in a string for safe JSON embedding.
func escapeJSONString(s string) string {
	var result strings.Builder
	for _, ch := range s {
		switch ch {
		case '"':
			result.WriteString(`\"`)
		case '\\':
			result.WriteString(`\\`)
		case '\n':
			result.WriteString(`\n`)
		case '\t':
			result.WriteString(`\t`)
		case '\r':
			result.WriteString(`\r`)
		default:
			result.WriteRune(ch)
		}
	}
	return result.String()
}