Novault-backend/internal/handler/ai_handler.go

package handler

import (
	"encoding/json"
	"io"
	"net/http"

	"accounting-app/internal/service"

	"github.com/gin-gonic/gin"
)

// AIHandler handles AI bookkeeping API requests
type AIHandler struct {
	aiService *service.AIBookkeepingService
}

// StreamChat handles streaming chat messages
// POST /api/v1/ai/chat/stream
func (h *AIHandler) StreamChat(c *gin.Context) {
	var req ChatRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"error":   "Invalid request: " + err.Error(),
		})
		return
	}

	// Get user ID from context
	userID := uint(1)
	if id, exists := c.Get("user_id"); exists {
		userID = id.(uint)
	}

	// Set headers for SSE
	c.Writer.Header().Set("Content-Type", "text/event-stream")
	c.Writer.Header().Set("Cache-Control", "no-cache")
	c.Writer.Header().Set("Connection", "keep-alive")
	c.Writer.Header().Set("Transfer-Encoding", "chunked")

	c.Stream(func(w io.Writer) bool {
		// Define the callback for chunks
		onChunk := func(chunk string) {
			// Sanitize chunk for SSE format (replace newlines to avoid breaking the stream protocol,
			// or just send raw if client handles it. Usually data: <content>\n\n)
			// For robustness, we JSON encode the data payload
			dataMap := map[string]string{"text": chunk}
			jsonData, _ := json.Marshal(dataMap)
			// message event
			c.SSEvent("message", string(jsonData))
		}

		// Call service
		response, err := h.aiService.StreamProcessChat(c.Request.Context(), userID, req.SessionID, req.Message, onChunk)

		if err != nil {
			// Send error event
			errMap := map[string]string{"error": err.Error()}
			jsonErr, _ := json.Marshal(errMap)
			c.SSEvent("error", string(jsonErr))
			return false
		}

		// Send final result event with metadata
		jsonResult, _ := json.Marshal(response)
		c.SSEvent("result", string(jsonResult))

		return false // Stop stream
	})
}

// NewAIHandler creates a new AIHandler
func NewAIHandler(aiService *service.AIBookkeepingService) *AIHandler {
	return &AIHandler{
		aiService: aiService,
	}
}

// ChatRequest represents a chat request
type ChatRequest struct {
	SessionID string `json:"session_id"`
	Message   string `json:"message" binding:"required"`
}

// InsightRequest represents an insight generation request
type InsightRequest struct {
	ContextData map[string]interface{} `json:"context_data" binding:"required"`
}

// TranscribeRequest represents a transcription request
type TranscribeRequest struct {
	// Audio file is sent as multipart form data
}

// ConfirmRequest represents a transaction confirmation request
type ConfirmRequest struct {
	SessionID string `json:"session_id" binding:"required"`
}

// RegisterRoutes registers AI routes
func (h *AIHandler) RegisterRoutes(rg *gin.RouterGroup) {
	ai := rg.Group("/ai")
	{
		ai.POST("/chat", h.Chat)
		ai.POST("/chat/stream", h.StreamChat) // New streaming endpoint
		ai.POST("/transcribe", h.Transcribe)
		ai.POST("/confirm", h.Confirm)
		ai.POST("/insight", h.Insight)
	}
}

// Chat handles chat messages for AI bookkeeping
// POST /api/v1/ai/chat
// Requirements: 12.1, 12.5
func (h *AIHandler) Chat(c *gin.Context) {
	var req ChatRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"error":   "Invalid request: " + err.Error(),
		})
		return
	}

	// Get user ID from context (default to 1 for now)
	userID := uint(1)
	if id, exists := c.Get("user_id"); exists {
		userID = id.(uint)
	}

	response, err := h.aiService.ProcessChat(c.Request.Context(), userID, req.SessionID, req.Message)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"error":   "Failed to process chat: " + err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"data":    response,
	})
}

// Insight handles daily insight generation
// POST /api/v1/ai/insight
func (h *AIHandler) Insight(c *gin.Context) {
	var req InsightRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"error":   "Invalid request: " + err.Error(),
		})
		return
	}

	userID := uint(1)
	if id, exists := c.Get("user_id"); exists {
		userID = id.(uint)
	}

	insightJSON, err := h.aiService.GenerateDailyInsight(c.Request.Context(), userID, req.ContextData)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"error":   "Failed to generate insight: " + err.Error(),
		})
		return
	}

	// Try to parse the JSON string from LLM to ensure structure
	var result map[string]interface{}
	if err := json.Unmarshal([]byte(insightJSON), &result); err != nil {
		// Fallback if LLM output isn't valid JSON (rare with correct prompt)
		result = map[string]interface{}{
			"spending": insightJSON,
			"budget":   "暂无建议",
		}
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"data":    result,
	})
}

// Transcribe handles audio transcription
// POST /api/v1/ai/transcribe
// Requirements: 12.2, 12.6
func (h *AIHandler) Transcribe(c *gin.Context) {
	// ... existing body ...
	file, header, err := c.Request.FormFile("audio")
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"error":   "No audio file provided: " + err.Error(),
		})
		return
	}
	defer file.Close()

	// Transcribe audio
	result, err := h.aiService.TranscribeAudio(c.Request.Context(), file, header.Filename)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"error":   "Failed to transcribe audio: " + err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"data":    result,
	})
}

// Confirm handles transaction confirmation
// POST /api/v1/ai/confirm
// Requirements: 12.3, 12.7, 12.8
func (h *AIHandler) Confirm(c *gin.Context) {
	var req ConfirmRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"success": false,
			"error":   "Invalid request: " + err.Error(),
		})
		return
	}

	// Get user ID from context (default to 1 for now)
	userID := uint(1)
	if id, exists := c.Get("user_id"); exists {
		userID = id.(uint)
	}

	transaction, err := h.aiService.ConfirmTransaction(c.Request.Context(), req.SessionID, userID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"success": false,
			"error":   "Failed to confirm transaction: " + err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"data":    transaction,
		"message": "Transaction created successfully",
	})
}