FranklinWH-Toolbox/service/controllers/aiChatControllers.go

package controllers

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/user"
	"path/filepath"
	"sort"
	"strings"
	"time"

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

	"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 模块
// runAsUser 指定拉起 kiro-cli 子进程时降权到的系统用户名，为空则不降权。
// home 指定 kiro-cli 子进程的 HOME 目录；为空时回退到降权用户的家目录，
// 若仍无法推导则使用默认值。
// debugLog 为调试日志宏开关，开启时向 kiro-cli 子进程注入 KIRO_LOG_LEVEL=debug。
func InitAIChat(serviceDir string, runAsUser string, home string, debugLog bool) {
	home = strings.TrimSpace(home)

	// 未显式配置 HOME 时的回退：优先推导降权用户的家目录，否则用默认值。
	if home == "" {
		home = "/home/jerry-remote"
		if u := strings.TrimSpace(runAsUser); u != "" {
			if usr, err := user.Lookup(u); err == nil {
				home = usr.HomeDir
			} else {
				zap.L().Warn("[AIChat] 查找降权用户失败，沿用默认 HOME",
					zap.String("user", u), zap.String("home", home), zap.Error(err))
			}
		}
	}

	// 组装传给 kiro-cli 子进程的环境变量。
	childEnv := []string{"HOME=" + home}
	if debugLog {
		// 宏开关：开启 kiro-cli debug 级日志。
		// 日志输出路径（Linux，kiro-cli acp 默认）：/tmp/kiro-log/logs/kiro-chat.log
		childEnv = append(childEnv, "KIRO_LOG_LEVEL=debug")
	}

	zap.L().Info("[AIChat] 初始化 ACP",
		zap.String("runAsUser", runAsUser), zap.String("home", home),
		zap.Bool("debugLog", debugLog))

	acpManager = services.NewACPManager(serviceDir, childEnv, runAsUser)
}

// 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) {
		// 连接生命周期上下文：写入失败（客户端断开）时取消，
		// 令 service/acp 层尽快停止本轮并解除阻塞。
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		// flush 封装：返回 false 表示连接已断开
		flush := func() bool {
			if err := w.Flush(); err != nil {
				cancel()
				return false
			}
			return true
		}

		// 发送 session_info 事件，包含实际生效的 agent 和 model
		infoContent, _ := json.Marshal(fiber.Map{"agentId": sess.AgentId, "modelId": sess.ModelId})
		infoEvent, _ := json.Marshal(services.AIChatEvent{Type: "session_info", Content: string(infoContent)})
		fmt.Fprintf(w, "event: session_info\ndata: %s\n\n", string(infoEvent))
		if !flush() {
			return
		}

		// 如果有欢迎消息且是新会话，先发送 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))
			if !flush() {
				return
			}
		}

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

		go acpManager.Prompt(ctx, 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))
			if !flush() {
				// 客户端断开：取消上下文，让 service 层停止转发并排空。
				// 继续 range 直到 eventCh 被关闭，避免 goroutine 泄漏。
				cancel()
				for range eventCh {
				}
				return
			}
		}
	})
	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)
}

// GetCurrentAiSession 获取当前活跃会话信息（含实际生效的 agent/model 与最近用量）
// GET /ai/sessions/current
func GetCurrentAiSession(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)
	sess := acpManager.GetSessionInfo(userId)
	if sess == nil {
		// 无活跃会话，返回空对象
		return Success(c, fiber.Map{})
	}

	var contextUsagePercent float64
	var turnDurationMs int64
	if sess.Client != nil {
		contextUsagePercent, turnDurationMs = sess.Client.GetMeta()
	}

	return Success(c, fiber.Map{
		"sessionId":           sess.SessionId,
		"title":               sess.Title,
		"agentId":             sess.AgentId,
		"modelId":             sess.ModelId,
		"contextUsagePercent": contextUsagePercent,
		"turnDurationMs":      turnDurationMs,
	})
}

// 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,
		"agentId":        sess.AgentId,
		"modelId":        sess.ModelId,
	})
}

// 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,
		"agentId":        sess.AgentId,
		"modelId":        sess.ModelId,
	})
}

// 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, "读取文件失败")
	}

	// 上传到 MinIO 的 "upload" 前缀（这样 Agent readTextFile 可以读到）
	record, err := filestore.UploadFile(userId, sess.SessionId, file.Filename, string(data))
	if err != nil {
		return Error(c, fmt.Sprintf("保存文件失败: %v", err))
	}

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

// ListAiFiles 获取用户的 AI 文件列表
// GET /ai/files/list?type=upload|download
func ListAiFiles(c *fiber.Ctx) error {
	userId := middleware.AuthGetUserName(c)
	fileType := utils.CopyString(c.Query("type"))

	if fileType != "upload" && fileType != "download" {
		return Fail(c, "type 参数必须为 upload 或 download")
	}

	// 按用户前缀列出 MinIO 文件
	prefix := fmt.Sprintf("%s/%s/", fileType, userId)
	objects, err := minioUtil.ListFilesObjectFromBucket(filestore.AIFileBucket, prefix, true)
	if err != nil {
		return Error(c, fmt.Sprintf("获取文件列表失败: %v", err))
	}

	type FileItem struct {
		Key       string `json:"key"`
		FileName  string `json:"fileName"`
		Size      int64  `json:"size"`
		CreatedAt int64  `json:"createdAt"` // Unix ms
		ExpiresAt int64  `json:"expiresAt"` // Unix ms（0 表示不过期）
	}

	now := time.Now()
	list := make([]FileItem, 0, len(objects))
	for _, obj := range objects {
		// 跳过已过期的文件（MinIO 可能延迟清理）
		if !obj.Expires.IsZero() && obj.Expires.Before(now) {
			continue
		}

		// 从 key 中提取文件名: {prefix}/{userId}/{sessionId}/{uuid}_{fileName}
		fileName := filepath.Base(obj.Key)
		if len(fileName) > 9 { // 去掉 "uuid_" 前缀（8位uuid + 下划线）
			fileName = fileName[9:]
		}

		var expiresAt int64
		if !obj.Expires.IsZero() {
			expiresAt = obj.Expires.UnixMilli()
		}

		list = append(list, FileItem{
			Key:       obj.Key,
			FileName:  fileName,
			Size:      obj.Size,
			CreatedAt: obj.LastModified.UnixMilli(),
			ExpiresAt: expiresAt,
		})
	}

	// 按创建时间倒序
	sort.Slice(list, func(i, j int) bool {
		return list[i].CreatedAt > list[j].CreatedAt
	})

	return Success(c, list)
}