digital-psychology/apps/api/internal/handler/ask.go

package handler

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"

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

	"github.com/yuxingu/digital-psychology/apps/api/internal/middleware"
	asksvc "github.com/yuxingu/digital-psychology/apps/api/internal/service/ask"
	"github.com/yuxingu/digital-psychology/apps/api/internal/textsafe"
	"github.com/yuxingu/digital-psychology/apps/api/pkg/response"
)

// AskHandler exposes AI 成长助手 APIs.
type AskHandler struct {
	Svc *asksvc.Service
}

// Register mounts ask routes.
func (h *AskHandler) Register(rg *gin.RouterGroup) {
	rg.GET("/ask/quota", h.GetQuota)
	rg.POST("/ask/threads", h.CreateThread)
	rg.DELETE("/ask/threads/:id", h.ClearThread)
	rg.GET("/ask/threads/:id/messages", h.ListMessages)
	rg.POST("/ask/threads/:id/messages", h.SendMessage)
}

// GetQuota handles GET /ask/quota.
func (h *AskHandler) GetQuota(c *gin.Context) {
	userID, ok := middleware.UserIDFromContext(c)
	if !ok {
		response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
		return
	}
	q, err := h.Svc.GetQuota(c.Request.Context(), userID)
	if err != nil {
		response.Fail(c, http.StatusInternalServerError, 50000, err.Error())
		return
	}
	response.OK(c, q)
}

// CreateThread handles POST /ask/threads.
func (h *AskHandler) CreateThread(c *gin.Context) {
	userID, ok := middleware.UserIDFromContext(c)
	if !ok {
		response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
		return
	}
	var req struct {
		ProfileID string `json:"profile_id" binding:"required"`
		Scene     string `json:"scene"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid request")
		return
	}
	pid, err := uuid.Parse(req.ProfileID)
	if err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid profile_id")
		return
	}
	th, err := h.Svc.CreateThread(c.Request.Context(), userID, asksvc.CreateThreadInput{
		ProfileID: pid, Scene: req.Scene,
	})
	if err != nil {
		if failTextCompliance(c, err) {
			return
		}
		response.Fail(c, http.StatusBadRequest, 40010, err.Error())
		return
	}
	response.OK(c, th)
}

// ClearThread handles DELETE /ask/threads/:id.
func (h *AskHandler) ClearThread(c *gin.Context) {
	userID, ok := middleware.UserIDFromContext(c)
	if !ok {
		response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
		return
	}
	tid, err := uuid.Parse(c.Param("id"))
	if err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid id")
		return
	}
	if err := h.Svc.ClearThread(c.Request.Context(), userID, tid); err != nil {
		response.Fail(c, http.StatusNotFound, 40410, err.Error())
		return
	}
	response.OK(c, gin.H{"cleared": true})
}

// ListMessages handles GET /ask/threads/:id/messages.
func (h *AskHandler) ListMessages(c *gin.Context) {
	userID, ok := middleware.UserIDFromContext(c)
	if !ok {
		response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
		return
	}
	tid, err := uuid.Parse(c.Param("id"))
	if err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid id")
		return
	}
	items, err := h.Svc.ListMessages(c.Request.Context(), userID, tid)
	if err != nil {
		response.Fail(c, http.StatusNotFound, 40410, err.Error())
		return
	}
	response.OK(c, gin.H{"items": items})
}

// SendMessage handles POST /ask/threads/:id/messages.
// Use ?stream=1 (or Accept: text/event-stream) for SSE streaming.
func (h *AskHandler) SendMessage(c *gin.Context) {
	userID, ok := middleware.UserIDFromContext(c)
	if !ok {
		response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
		return
	}
	tid, err := uuid.Parse(c.Param("id"))
	if err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid id")
		return
	}
	var req struct {
		Content string `json:"content" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		response.Fail(c, http.StatusBadRequest, 10000, "invalid request")
		return
	}

	wantStream := c.Query("stream") == "1" ||
		strings.Contains(c.GetHeader("Accept"), "text/event-stream")
	if wantStream {
		h.sendMessageStream(c, userID, tid, req.Content)
		return
	}

	out, err := h.Svc.SendMessage(c.Request.Context(), userID, tid, req.Content)
	if err != nil {
		if failTextCompliance(c, err) {
			return
		}
		if asksvc.IsQuotaExhausted(err) {
			response.Fail(c, http.StatusPaymentRequired, 40210, "问答次数已用完，可购买额度或开通成长会员")
			return
		}
		response.Fail(c, http.StatusBadRequest, 40011, err.Error())
		return
	}
	response.OK(c, out)
}

func (h *AskHandler) sendMessageStream(c *gin.Context, userID, tid uuid.UUID, content string) {
	c.Writer.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	c.Writer.Header().Set("Cache-Control", "no-cache, no-transform")
	c.Writer.Header().Set("Connection", "keep-alive")
	c.Writer.Header().Set("X-Accel-Buffering", "no")
	c.Status(http.StatusOK)
	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		response.Fail(c, http.StatusInternalServerError, 50000, "stream unsupported")
		return
	}

	writeEvent := func(event string, payload any) error {
		raw, err := json.Marshal(payload)
		if err != nil {
			return err
		}
		if _, err := fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, raw); err != nil {
			return err
		}
		flusher.Flush()
		return nil
	}

	err := h.Svc.StreamMessage(c.Request.Context(), userID, tid, content, writeEvent)
	if err != nil {
		msg := err.Error()
		code := 40011
		if errors.Is(err, textsafe.ErrRejected) {
			code = 40060
		} else if asksvc.IsQuotaExhausted(err) {
			msg = "问答次数已用完，可购买额度或开通成长会员"
			code = 40210
		}
		_ = writeEvent("error", map[string]any{"code": code, "message": msg})
	}
}