PR-Helper/handlers/generate.go

package handlers

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/HoHD/PR-Helper/services"
	"github.com/gin-gonic/gin"
)

type GenerateHandler struct {
	db *sql.DB
}

func NewGenerateHandler(db *sql.DB) *GenerateHandler {
	return &GenerateHandler{db: db}
}

// Generate handles POST /api/repos/:id/generate — SSE streaming PR description generation.
func (h *GenerateHandler) Generate(c *gin.Context) {
	user := GetCurrentUser(c)
	if user == nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "未登录"})
		return
	}

	id := c.Param("id")

	// Get repo info (scoped to user)
	var localPath string
	err := h.db.QueryRow(`SELECT local_path FROM repositories WHERE id = ? AND user_id = ?`, id, user.ID).Scan(&localPath)
	if err == sql.ErrNoRows {
		c.JSON(http.StatusNotFound, gin.H{"error": "仓库未找到"})
		return
	}
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// Parse request
	var req struct {
		Base string `json:"base" binding:"required"`
		Head string `json:"head" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "请选择 Base 和 Head 分支"})
		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")
	c.Status(http.StatusOK)

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器不支持流式传输"})
		return
	}

	sendEvent := func(event string, data interface{}) {
		jsonData, err := json.Marshal(data)
		if err != nil {
			jsonData = []byte(`{"error":"序列化事件数据失败"}`)
		}
		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, jsonData)
		flusher.Flush()
	}

	// Update last_used
	h.db.Exec(`UPDATE repositories SET last_used = NOW() WHERE id = ?`, id)

	// Generate PR description (streams markdown via content events)
	markdown, err := services.GeneratePR(h.db, localPath, req.Base, req.Head, user.ID, sendEvent)
	if err != nil {
		sendEvent("error", map[string]interface{}{"message": err.Error()})
		return
	}

	// Save analysis to DB
	resultJSON, _ := json.Marshal(map[string]string{"markdown": markdown})
	if _, err := h.db.Exec(`INSERT INTO analyses (user_id, repo_id, type, base_ref, head_ref, result) VALUES (?, ?, 'pr_description', ?, ?, ?)`,
		user.ID, id, req.Base, req.Head, string(resultJSON)); err != nil {
		sendEvent("error", map[string]interface{}{"message": "保存分析结果失败: " + err.Error()})
		return
	}
}