FranklinWH-Toolbox/service/controllers/aiChatControllers.go

package controllers

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"path/filepath"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/utils"
	"github.com/minio/minio-go/v7"

	"fwh-toolbox-server/middleware"
	"fwh-toolbox-server/services"
	"fwh-toolbox-server/utils/filestore"
	minioUtil "fwh-toolbox-server/utils/minio"
)

var acpManager *services.ACPManager

// InitAIChat 初始化 AI Chat 模块
func InitAIChat(serviceDir string) {
	acpManager = services.NewACPManager(serviceDir, []string{"HOME=/home/jerry-remote"})
}

// AiChat 发送消息（SSE 流式响应）
// POST /ai/chat
func AiChat(c *fiber.Ctx) error {
	var req struct {
		Message   string `json:"message"`
		SessionId string `json:"sessionId"`
	}
	if err := c.BodyParser(&req); err != nil {
		return Fail(c, "Invalid request body")
	}
	if req.Message == "" {
		return Fail(c, "message 不能为空")
	}

	userId := middleware.AuthGetUserName(c)

	sess, err := acpManager.GetOrCreateSession(userId)
	if err != nil {
		return Error(c, fmt.Sprintf("AI 会话创建失败: %v", err))
	}

	// SSE 流式响应
	c.Set("Content-Type", "text/event-stream")
	c.Set("Cache-Control", "no-cache")
	c.Set("Connection", "keep-alive")
	c.Set("X-Accel-Buffering", "no")

	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
		// 如果有欢迎消息且是新会话，先发送 welcome 事件
		if sess.WelcomeMessage != "" {
			welcomeEvent, _ := json.Marshal(services.AIChatEvent{Type: "welcome", Content: sess.WelcomeMessage})
			fmt.Fprintf(w, "event: welcome\ndata: %s\n\n", string(welcomeEvent))
			w.Flush()
		}

		eventCh := make(chan services.AIChatEvent, 50)

		go acpManager.Prompt(sess, req.Message, eventCh)

		for event := range eventCh {
			data, _ := json.Marshal(event)
			fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.Type, string(data))
			w.Flush()
		}
	})
	return nil
}

// GetAiSessionList 获取用户的会话历史列表
// GET /ai/sessions/list
func GetAiSessionList(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)
	list := acpManager.GetSessionList(userId)
	return Success(c, list)
}

// CreateAiSession 创建新的 AI 会话（保留旧会话到历史）
// POST /ai/sessions/new
func CreateAiSession(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)

	// 创建新会话（旧会话自动归档到历史）
	sess, err := acpManager.CreateNewSession(userId)
	if err != nil {
		return Error(c, fmt.Sprintf("创建会话失败: %v", err))
	}

	return Success(c, fiber.Map{
		"sessionId":      sess.SessionId,
		"welcomeMessage": sess.WelcomeMessage,
	})
}

// SwitchAiSession 切换到历史会话（恢复 AI 上下文，轻量响应）
// POST /ai/sessions/:id
func SwitchAiSession(c *fiber.Ctx) error {
	sessionId := utils.CopyString(c.Params("id"))
	if sessionId == "" {
		return Fail(c, "sessionId 不能为空")
	}

	userId := middleware.AuthGetUserName(c)
	sess, err := acpManager.LoadSession(userId, sessionId)
	if err != nil {
		return Error(c, fmt.Sprintf("加载会话失败: %v", err))
	}

	return Success(c, fiber.Map{
		"sessionId":      sess.SessionId,
		"title":          sess.Title,
		"welcomeMessage": sess.WelcomeMessage,
	})
}

// StreamLoadAiSession SSE 流式推送回放消息
// GET /ai/sessions/:id/load
func StreamLoadAiSession(c *fiber.Ctx) error {
	sessionId := utils.CopyString(c.Params("id"))
	if sessionId == "" {
		return Fail(c, "sessionId 不能为空")
	}

	userId := middleware.AuthGetUserName(c)

	// 获取回放消息（会话必须已通过 SwitchAiSession 加载）
	replayMessages := acpManager.GetReplayMessages(userId)
	if replayMessages == nil {
		return Fail(c, "会话未加载，请先调用切换会话接口")
	}

	// SSE 流式响应
	c.Set("Content-Type", "text/event-stream")
	c.Set("Cache-Control", "no-cache")
	c.Set("Connection", "keep-alive")
	c.Set("X-Accel-Buffering", "no")

	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
		var totalBytes int

		for _, msg := range replayMessages {
			data, _ := json.Marshal(msg)
			totalBytes += len(msg.Content)
			fmt.Fprintf(w, "event: message\ndata: %s\n\n", string(data))
			w.Flush()
		}

		// 发送 done 事件，包含统计信息和建议
		suggestNewSession := totalBytes >= 100*1024 // 100KB 以上建议新建
		donePayload, _ := json.Marshal(fiber.Map{
			"sessionId":         sessionId,
			"totalMessages":     len(replayMessages),
			"totalBytes":        totalBytes,
			"suggestNewSession": suggestNewSession,
		})
		fmt.Fprintf(w, "event: done\ndata: %s\n\n", string(donePayload))
		w.Flush()
	})
	return nil
}

// UpdateAiSessionTitle 更新会话标题
// PATCH /ai/sessions/:id
func UpdateAiSessionTitle(c *fiber.Ctx) error {
	sessionId := utils.CopyString(c.Params("id"))
	if sessionId == "" {
		return Fail(c, "sessionId 不能为空")
	}

	var req struct {
		Title string `json:"title"`
	}
	if err := c.BodyParser(&req); err != nil {
		return Fail(c, "Invalid request body")
	}

	userId := middleware.AuthGetUserName(c)
	acpManager.UpdateSessionTitle(userId, sessionId, req.Title)
	return Success[any](c, nil, "标题已更新")
}

// DeleteAiSession 删除 AI 会话
// DELETE /ai/sessions/:id
func DeleteAiSession(c *fiber.Ctx) error {
	sessionId := utils.CopyString(c.Params("id"))
	userId := middleware.AuthGetUserName(c)

	if sessionId == "" {
		// 无 id 参数，删除当前活跃会话
		acpManager.DeleteSession(userId)
	} else {
		// 删除指定会话
		acpManager.DeleteSessionFromHistory(userId, sessionId)
	}

	return Success[any](c, nil, "会话已删除")
}

// CancelAiChat 取消当前 AI 操作
// POST /ai/cancel
func CancelAiChat(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)
	sess := acpManager.GetSessionInfo(userId)
	if sess == nil {
		return Fail(c, "没有活跃的会话")
	}
	acpManager.Cancel(sess)
	return Success[any](c, nil, "已取消")
}

// ResolvePermission 前端回传写权限选择
// POST /ai/permission
func ResolvePermission(c *fiber.Ctx) error {
	var req struct {
		OptionId string `json:"optionId"`
	}
	if err := c.BodyParser(&req); err != nil {
		return Fail(c, "Invalid request body")
	}
	if req.OptionId == "" {
		return Fail(c, "optionId 不能为空")
	}

	userId := middleware.AuthGetUserName(c)
	if err := acpManager.ResolvePermission(userId, req.OptionId); err != nil {
		return Fail(c, err.Error())
	}

	return Success[any](c, nil, "已回复")
}

// DownloadAiFile 下载 AI 生成的文件
// GET /ai/files/download?key=xxx
func DownloadAiFile(c *fiber.Ctx) error {
	key := utils.CopyString(c.Query("key"))
	if key == "" {
		return Fail(c, "key 不能为空")
	}

	// 直接从 MinIO 读取文件（key 即 objectKey）
	obj, err := minioUtil.Client.GetObject(context.Background(), filestore.AIFileBucket, key, minio.GetObjectOptions{})
	if err != nil {
		return Fail(c, "文件不存在")
	}
	defer obj.Close()

	// 获取文件信息
	stat, err := obj.Stat()
	if err != nil {
		return Fail(c, fmt.Sprintf("文件不存在: %v", err))
	}

	// 读取全部内容
	data, err := io.ReadAll(obj)
	if err != nil {
		return Error(c, fmt.Sprintf("读取文件失败: %v", err))
	}

	// 从 key 中提取文件名（key 格式: prefix/userId/sessionId/uuid_fileName）
	fileName := filepath.Base(key)
	// 去掉 uuid 前缀（前9个字符: 8位uuid + 下划线）
	if len(fileName) > 9 {
		fileName = fileName[9:]
	}

	c.Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, fileName))
	c.Set("Content-Type", stat.ContentType)
	return c.Send(data)
}

// UploadAiFile 上传文件给 AI 分析
// POST /ai/files/upload
func UploadAiFile(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)

	sess, err := acpManager.GetOrCreateSession(userId)
	if err != nil {
		return Error(c, fmt.Sprintf("AI 会话创建失败: %v", err))
	}

	// 读取上传文件
	file, err := c.FormFile("file")
	if err != nil {
		return Fail(c, "请选择要上传的文件")
	}

	// 限制文件大小（最大 2MB）
	if file.Size > 2*1024*1024 {
		return Fail(c, "文件大小不能超过 2MB")
	}

	// 读取文件内容
	f, err := file.Open()
	if err != nil {
		return Error(c, "打开文件失败")
	}
	defer f.Close()

	data, err := io.ReadAll(f)
	if err != nil {
		return Error(c, "读取文件失败")
	}

	// 写入 FileStore（存到 ACP 的 FileStore，这样 Agent readTextFile 可以读到）
	virtualPath := fmt.Sprintf("/uploads/%s", file.Filename)
	record, err := sess.Client.FileStore.WriteFile(userId, sess.SessionId, virtualPath, string(data))
	if err != nil {
		return Error(c, fmt.Sprintf("保存文件失败: %v", err))
	}

	return Success(c, fiber.Map{
		"key":      record.Key,
		"fileName": record.FileName,
		"path":     record.Path,
		"size":     record.Size,
	})
}