FranklinWH-Toolbox/service/services/aiChatServices.go

package services

import (
	"context"
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/repositories"
	"fwh-toolbox-server/utils/acp"
	"os"
	"os/exec"
	"sync"
	"time"

	"go.uber.org/zap"
)

// AIChatEvent SSE 事件（输出给前端）
type AIChatEvent struct {
	Type        string `json:"type"`                  // text, tool_call, tool_update, done, error
	Content     string `json:"content,omitempty"`     // 文本内容
	ToolCallId  string `json:"toolCallId,omitempty"`  // 工具调用唯一ID（前端按 id 关联状态）
	ToolName    string `json:"name,omitempty"`        // 工具原始 title（前端作 hover tooltip）
	Description string `json:"description,omitempty"` // 工具简要描述（前端展示）
	Status      string `json:"status,omitempty"`      // 工具状态（已归一化：running/completed/error）
	SessionId   string `json:"sessionId,omitempty"`   // 会话ID
	StopReason  string `json:"stopReason,omitempty"`
}

// ACPSession 单个用户的 ACP 会话
type ACPSession struct {
	UserId         string
	SessionId      string
	Title          string
	WelcomeMessage string
	AgentId        string              // 实际生效的 agent（currentModeId）
	ModelId        string              // 实际生效的 model（currentModelId）
	ReplayMessages []acp.ReplayMessage // session/load 回放的历史消息
	Client         *acp.Client
	Mutex          sync.Mutex
	LastActive     time.Time
	CreatedAt      time.Time
}

// SessionListItem 会话列表项（返回给前端）
type SessionListItem struct {
	SessionId  string `json:"sessionId"`
	Title      string `json:"title"`
	CreatedAt  int64  `json:"createdAt"`
	LastActive int64  `json:"lastActive"`
	IsActive   bool   `json:"isActive"`
}

// ACPManager 管理所有用户的 ACP 会话
type ACPManager struct {
	sessions   map[string]*ACPSession // key: userId -> 当前活跃会话（持有进程引用）
	mu         sync.Mutex
	serviceDir string
	acpConfig  acp.ClientConfig
}

// NewACPManager 创建 ACP 管理器
// runAsUser 指定拉起 kiro-cli 子进程时降权到的系统用户名，为空则不降权。
func NewACPManager(serviceDir string, env []string, runAsUser string) *ACPManager {
	return &ACPManager{
		sessions:   make(map[string]*ACPSession),
		serviceDir: serviceDir,
		acpConfig: acp.ClientConfig{
			Command:    "kiro-cli",
			Args:       []string{"acp", "--agent", "fwh-toolbox-agent"},
			WorkingDir: serviceDir,
			Env:        env,
			RunAsUser:  runAsUser,
		},
	}
}

// ============ 会话生命周期 ============

// GetOrCreateSession 获取或创建用户的 ACP 会话
// 优先级：内存活跃会话 → Redis 持久化恢复 → 创建新会话
func (m *ACPManager) GetOrCreateSession(userId string) (*ACPSession, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 1. 检查内存中是否有活跃会话
	if sess, ok := m.sessions[userId]; ok {
		if sess.Client.IsAlive() {
			sess.LastActive = time.Now()
			return sess, nil
		}
		sess.Client.Close()
		delete(m.sessions, userId)
	}

	// 2. 尝试从 Redis 恢复会话
	savedSession, err := repositories.GetAiSession(userId)
	if err == nil && savedSession != nil {
		sess, err := m.startAndLoadSession(userId, savedSession)
		if err == nil {
			zap.L().Info("[AIChat] 从 Redis 恢复会话成功", zap.String("userId", userId), zap.String("sessionId", savedSession.SessionId))
			return sess, nil
		}
		zap.L().Warn("[AIChat] 从 Redis 恢复会话失败，将创建新会话", zap.String("userId", userId), zap.Error(err))
		repositories.DeleteAiSession(userId)
	}

	// 3. 创建新会话
	return m.createSession(userId)
}

// CreateNewSession 强制创建一个新会话（归档旧会话到 Redis 历史）
func (m *ACPManager) CreateNewSession(userId string) (*ACPSession, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 归档当前活跃会话
	if sess, ok := m.sessions[userId]; ok {
		m.archiveSession(userId, sess)
		sess.Client.Close()
		delete(m.sessions, userId)
	}

	return m.createSession(userId)
}

// LoadSession 加载历史会话（恢复 AI 上下文，切换会话）
// 仅切换会话并恢复上下文，不返回回放消息（回放消息通过 SSE 接口单独获取）
func (m *ACPManager) LoadSession(userId string, sessionId string) (*ACPSession, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 如果已经是当前会话，直接返回
	if sess, ok := m.sessions[userId]; ok {
		if sess.SessionId == sessionId {
			sess.LastActive = time.Now()
			return sess, nil
		}
		// 归档当前活跃会话
		m.archiveSession(userId, sess)
		sess.Client.Close()
		delete(m.sessions, userId)
	}

	// 从 Redis 历史列表中获取标题
	item := m.getSessionListFromRedisHistory(userId, sessionId)
	if item == nil {
		return nil, fmt.Errorf("redis session/load 失败")
	}

	// 从 Redis 历史列表中移除该会话（因为它即将变为活跃）
	m.removeFromRedisHistory(userId, sessionId)

	// 启动 ACP 进程并 load 会话
	client, err := acp.NewClient(m.acpConfig)
	if err != nil {
		return nil, err
	}
	if err := client.Initialize(); err != nil {
		client.Close()
		return nil, err
	}

	loadedId, welcomeMsg, agentId, modelId, replayMessages, err := client.LoadSession(sessionId, m.serviceDir)
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("session/load 失败: %w", err)
	}

	zap.L().Info("[AIChat] session/load 回放完成",
		zap.String("sessionId", loadedId),
		zap.String("welcomeMsg", welcomeMsg),
		zap.Int("replayCount", len(replayMessages)),
	)

	sess := &ACPSession{
		UserId:         userId,
		SessionId:      sessionId,
		Title:          item.Title,
		WelcomeMessage: welcomeMsg,
		AgentId:        agentId,
		ModelId:        modelId,
		ReplayMessages: replayMessages,
		Client:         client,
		LastActive:     time.Now(),
		CreatedAt:      time.UnixMilli(item.CreatedAt),
	}
	client.UserId = userId
	m.sessions[userId] = sess

	// 持久化
	m.activateSession(userId, sess)

	zap.L().Info("[AIChat] 会话加载成功", zap.String("userId", userId), zap.String("sessionId", sessionId))
	return sess, nil
}

// GetReplayMessages 获取指定用户当前活跃会话的回放消息
// 用于 SSE 流式推送
func (m *ACPManager) GetReplayMessages(userId string) []acp.ReplayMessage {
	m.mu.Lock()
	defer m.mu.Unlock()

	if sess, ok := m.sessions[userId]; ok {
		return sess.ReplayMessages
	}
	return nil
}

// ============ 会话列表查询 ============

// GetSessionList 获取用户的会话列表（当前活跃 + Redis 历史）
func (m *ACPManager) GetSessionList(userId string) []SessionListItem {
	m.mu.Lock()
	defer m.mu.Unlock()

	list := make([]SessionListItem, 0)

	// 当前活跃会话排在最前
	if sess, ok := m.sessions[userId]; ok {
		list = append(list, SessionListItem{
			SessionId:  sess.SessionId,
			Title:      sess.Title,
			CreatedAt:  sess.CreatedAt.UnixMilli(),
			LastActive: sess.LastActive.UnixMilli(),
			IsActive:   true,
		})
	} else {
		// 内存无活跃会话（后端重启），从 Redis 恢复完整会话到内存
		activeSess, err := repositories.GetAiSession(userId)
		if err == nil && activeSess != nil {
			sess, err := m.startAndLoadSession(userId, activeSess)
			if err == nil {
				list = append(list, SessionListItem{
					SessionId:  sess.SessionId,
					Title:      sess.Title,
					CreatedAt:  sess.CreatedAt.UnixMilli(),
					LastActive: sess.LastActive.UnixMilli(),
					IsActive:   true,
				})
				zap.L().Info("[AIChat] GetSessionList 恢复会话到内存", zap.String("userId", userId), zap.String("sessionId", sess.SessionId))
			}
		}
	}

	// 从 Redis 获取历史列表
	redisItems, err := repositories.GetAiSessionList(userId)
	if err == nil && len(redisItems) > 0 {
		for _, item := range redisItems {
			list = append(list, SessionListItem{
				SessionId:  item.SessionId,
				Title:      item.Title,
				CreatedAt:  item.CreatedAt,
				LastActive: item.LastActive,
				IsActive:   false,
			})
		}
	}

	return list
}

// ============ 会话操作 ============

// UpdateSessionTitle 更新会话标题
func (m *ACPManager) UpdateSessionTitle(userId string, sessionId string, title string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 检查当前活跃会话
	if sess, ok := m.sessions[userId]; ok && sess.SessionId == sessionId {
		sess.Title = title
		m.saveSessionTitle(userId, sess)
		return
	}

	// 历史会话
	m.saveHistorySessionTitle(userId, sessionId, title)
}

// DeleteSessionFromHistory 删除某会话
func (m *ACPManager) DeleteSessionFromHistory(userId string, sessionId string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 如果是当前活跃会话，关闭它
	if sess, ok := m.sessions[userId]; ok && sess.SessionId == sessionId {
		sess.Client.Close()
		delete(m.sessions, userId)
		m.removeSession(userId, sessionId, true)
		return
	}

	// 历史会话
	m.removeSession(userId, sessionId, false)
}

// DeleteSession 删除当前活跃会话
func (m *ACPManager) DeleteSession(userId string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if sess, ok := m.sessions[userId]; ok {
		sess.Client.Close()
		delete(m.sessions, userId)
		m.removeSession(userId, sess.SessionId, true)
		zap.L().Info("[AIChat] 会话已删除", zap.String("userId", userId))
	}
}

// GetSessionInfo 获取用户当前活跃会话信息
func (m *ACPManager) GetSessionInfo(userId string) *ACPSession {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.sessions[userId]
}

// Cancel 取消当前操作
func (m *ACPManager) Cancel(sess *ACPSession) {
	sess.Client.Cancel(sess.SessionId)
}

// ResolvePermission 前端回传权限选择
func (m *ACPManager) ResolvePermission(userId string, optionId string) error {
	sess := m.GetSessionInfo(userId)
	if sess == nil {
		return fmt.Errorf("没有活跃的会话")
	}
	sess.Client.ResolvePermission(optionId)
	return nil
}

// ============ 消息发送 ============

// Prompt 发送消息并流式读取响应。
// ctx 由上层（HTTP 处理器）在连接断开时取消：一旦取消，立即通知 Agent 取消本轮
// （Client.Cancel）、唤醒 acp 层阻塞的事件发送（Client.AbortEvents），并停止向
// eventCh 转发，避免 goroutine 因客户端断开而永久阻塞。
func (m *ACPManager) Prompt(ctx context.Context, sess *ACPSession, message string, eventCh chan<- AIChatEvent) {
	sess.Mutex.Lock()
	defer sess.Mutex.Unlock()

	sess.LastActive = time.Now()

	// 首次发消息时自动设置标题（取消息前 30 字）
	if sess.Title == "" && message != "" {
		title := message
		if len([]rune(title)) > 30 {
			title = string([]rune(title)[:30]) + "..."
		}
		sess.Title = title
	}

	// 注入会话上下文（让 AI 在调用 MCP Tool 时能传入正确的 userId 和 sessionId）
	contextMessage := fmt.Sprintf("[context] userId=%s, sessionId=%s\n\n%s", sess.UserId, sess.SessionId, message)

	// 用 acp.Client.Prompt 流式读取
	acpEvents := make(chan acp.Event, 50)
	go sess.Client.Prompt(sess.SessionId, contextMessage, acpEvents)

	// 监听上层取消：断开时通知 Agent 取消并解除 acp 层发送阻塞
	go func() {
		<-ctx.Done()
		sess.Client.Cancel(sess.SessionId)
		sess.Client.AbortEvents()
	}()

	// 转发单条事件，支持在上层取消时中断，避免阻塞在满通道上
	send := func(ev AIChatEvent) bool {
		select {
		case eventCh <- ev:
			return true
		case <-ctx.Done():
			return false
		}
	}

	// 转换 acp.Event → AIChatEvent
	for event := range acpEvents {
		var ok bool
		switch event.Type {
		case "text":
			ok = send(AIChatEvent{Type: "text", Content: event.Content})
		case "tool_call":
			ok = send(AIChatEvent{
				Type:        "tool_call",
				ToolCallId:  event.ToolCallId,
				ToolName:    event.ToolName,
				Description: acp.DescribeTool(event.ToolName),
				Status:      acp.NormalizeStatus(event.ToolStatus),
			})
		case "tool_update":
			// update 一般只变更状态；若本次携带了新 title 才重新映射简述
			ev := AIChatEvent{
				Type:       "tool_update",
				ToolCallId: event.ToolCallId,
				Status:     acp.NormalizeStatus(event.ToolStatus),
			}
			if event.ToolName != "" {
				ev.ToolName = event.ToolName
				ev.Description = acp.DescribeTool(event.ToolName)
			}
			ok = send(ev)
		case "permission":
			ok = send(AIChatEvent{Type: "permission", Content: toJSON(event.Permission)})
		case "metadata":
			ok = send(AIChatEvent{Type: "metadata", Content: event.Content})
		case "done":
			ok = send(AIChatEvent{Type: "done", SessionId: event.SessionId, StopReason: event.StopReason})
		case "error":
			ok = send(AIChatEvent{Type: "error", Content: event.Content})
		default:
			ok = true
		}
		if !ok {
			// 上层已取消：继续排空 acpEvents 直至 acp 层因 abort/取消关闭它，
			// 但不再向 eventCh 转发，防止阻塞。
			for range acpEvents {
			}
			break
		}
	}

	close(eventCh)
}

// ============ 内部方法 ============

// createSession 创建新 ACP 会话
func (m *ACPManager) createSession(userId string) (*ACPSession, error) {
	client, err := acp.NewClient(m.acpConfig)
	if err != nil {
		return nil, err
	}
	if err := client.Initialize(); err != nil {
		client.Close()
		return nil, err
	}

	sessionId, welcomeMsg, agentId, modelId, err := client.NewSession(m.serviceDir)
	if err != nil {
		client.Close()
		return nil, err
	}

	now := time.Now()
	sess := &ACPSession{
		UserId:         userId,
		SessionId:      sessionId,
		Title:          "",
		WelcomeMessage: welcomeMsg,
		AgentId:        agentId,
		ModelId:        modelId,
		Client:         client,
		LastActive:     now,
		CreatedAt:      now,
	}
	client.UserId = userId
	m.sessions[userId] = sess

	// 持久化
	m.saveNewSession(userId, sess)

	zap.L().Info("[AIChat] 会话创建成功", zap.String("userId", userId), zap.String("sessionId", sessionId), zap.String("agent", agentId), zap.String("model", modelId))
	return sess, nil
}

// startAndLoadSession 启动 ACP 进程并 load 指定会话（用于 Redis 恢复）
func (m *ACPManager) startAndLoadSession(userId string, active *repositories.AiActiveSession) (*ACPSession, error) {
	client, err := acp.NewClient(m.acpConfig)
	if err != nil {
		return nil, err
	}
	if err := client.Initialize(); err != nil {
		client.Close()
		return nil, err
	}

	_, welcomeMsg, agentId, modelId, _, err := client.LoadSession(active.SessionId, m.serviceDir)
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("session/load 失败: %w", err)
	}

	now := time.Now()
	sess := &ACPSession{
		UserId:         userId,
		SessionId:      active.SessionId,
		Title:          active.Title,
		WelcomeMessage: welcomeMsg,
		AgentId:        agentId,
		ModelId:        modelId,
		Client:         client,
		LastActive:     now,
		CreatedAt:      time.UnixMilli(active.CreatedAt),
	}
	client.UserId = userId
	m.sessions[userId] = sess
	return sess, nil
}

// ============ 内部方法：会话存储 ============

// persistNewSession 持久化新创建的活跃会话（Redis + MongoDB）
func (m *ACPManager) saveNewSession(userId string, sess *ACPSession) {
	repositories.SaveAiSession(userId, repositories.AiActiveSession{
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		CreatedAt:  sess.CreatedAt.UnixMilli(),
		LastActive: sess.LastActive.UnixMilli(),
	})
	repositories.DeactivateAllAiChatSessions(userId)
	repositories.UpsertAiChatSession(repositories.AiChatSession{
		UserId:     userId,
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		IsActive:   true,
		CreatedAt:  sess.CreatedAt,
		LastActive: sess.LastActive,
	})
}

// persistActivateSession 持久化会话激活（Redis + MongoDB）
func (m *ACPManager) activateSession(userId string, sess *ACPSession) {
	repositories.SaveAiSession(userId, repositories.AiActiveSession{
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		CreatedAt:  sess.CreatedAt.UnixMilli(),
		LastActive: sess.LastActive.UnixMilli(),
	})
	repositories.DeactivateAllAiChatSessions(userId)
	repositories.UpdateAiChatSessionActive(userId, sess.SessionId, true)
}

// persistArchiveSession 持久化会话归档（Redis + MongoDB）
func (m *ACPManager) archiveSession(userId string, sess *ACPSession) {
	item := repositories.AiSessionItem{
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		CreatedAt:  sess.CreatedAt.UnixMilli(),
		LastActive: sess.LastActive.UnixMilli(),
	}

	list, _ := repositories.GetAiSessionList(userId)
	list = append([]repositories.AiSessionItem{item}, list...)

	if len(list) > 50 {
		overflow := list[50:]
		list = list[:50]
		go func() {
			for _, old := range overflow {
				m.deleteKiroSession(old.SessionId)
			}
		}()
	}

	if err := repositories.SaveAiSessionList(userId, list); err != nil {
		zap.L().Warn("[AIChat] 归档会话到Redis失败", zap.String("userId", userId), zap.Error(err))
	}
	repositories.UpdateAiChatSessionActive(userId, sess.SessionId, false)
}

// persistUpdateTitle 持久化标题更新（Redis + MongoDB）
func (m *ACPManager) saveSessionTitle(userId string, sess *ACPSession) {
	repositories.SaveAiSession(userId, repositories.AiActiveSession{
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		CreatedAt:  sess.CreatedAt.UnixMilli(),
		LastActive: sess.LastActive.UnixMilli(),
	})
	repositories.UpdateAiChatSessionTitle(userId, sess.SessionId, sess.Title)
}

// persistUpdateHistoryTitle 持久化历史会话标题更新（Redis + MongoDB）
func (m *ACPManager) saveHistorySessionTitle(userId string, sessionId string, title string) {
	repositories.UpdateAiSessionTitle(userId, sessionId, title)
	repositories.UpdateAiChatSessionTitle(userId, sessionId, title)
}

// persistDeleteSession 持久化删除会话（Redis + MongoDB）
func (m *ACPManager) removeSession(userId string, sessionId string, isActive bool) {
	if isActive {
		repositories.DeleteAiSession(userId)
	} else {
		m.removeFromRedisHistory(userId, sessionId)
	}
	repositories.DeleteAiChatSession(userId, sessionId)
}

// ============ 内部方法：Redis 操作 ============
func (m *ACPManager) removeFromRedisHistory(userId string, sessionId string) {
	list, err := repositories.GetAiSessionList(userId)
	if err != nil || len(list) == 0 {
		return
	}

	for i := range list {
		if list[i].SessionId == sessionId {
			list = append(list[:i], list[i+1:]...)
			repositories.SaveAiSessionList(userId, list)
			return
		}
	}
}

// getSessionListFromRedisHistory 从 Redis 历史列表中获取会话
func (m *ACPManager) getSessionListFromRedisHistory(userId string, sessionId string) *repositories.AiSessionItem {
	list, err := repositories.GetAiSessionList(userId)
	if err != nil {
		zap.L().Warn("[AIChat] 从redis读取AiSessionList失败", zap.String("err", err.Error()))
		return nil
	}
	for _, item := range list {
		if item.SessionId == sessionId {
			return &item
		}
	}
	return nil
}

// deleteKiroSession 通过 kiro-cli 命令删除会话
func (m *ACPManager) deleteKiroSession(sessionId string) {
	cmd := exec.Command("kiro-cli", "chat", "--delete-session", sessionId)
	cmd.Dir = m.serviceDir
	// 复用与 ACP 拉起一致的环境变量（如 HOME），保证操作的是同一份会话数据
	cmd.Env = append(os.Environ(), m.acpConfig.Env...)
	// 与拉起子进程一致地降权，避免以 root 身份写入目标用户的会话文件
	if err := acp.ApplyRunAsUser(cmd, m.acpConfig.RunAsUser); err != nil {
		zap.L().Warn("[AIChat] 删除Kiro会话降权失败，跳过删除",
			zap.String("sessionId", sessionId),
			zap.Error(err),
		)
		return
	}
	if output, err := cmd.CombinedOutput(); err != nil {
		zap.L().Warn("[AIChat] 删除Kiro会话失败",
			zap.String("sessionId", sessionId),
			zap.Error(err),
			zap.String("output", string(output)),
		)
	} else {
		zap.L().Info("[AIChat] 删除Kiro会话成功", zap.String("sessionId", sessionId))
	}
}

// ============ 工具函数 ============

func toJSON(v any) string {
	data, _ := json.Marshal(v)
	return string(data)
}