FranklinWH-Toolbox/service/services/aiChatServices.go

package services

import (
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/utils/acp"
	"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"`   // 文本内容
	ToolName   string `json:"name,omitempty"`      // 工具名
	Status     string `json:"status,omitempty"`    // 工具状态
	SessionId  string `json:"sessionId,omitempty"` // 会话ID
	StopReason string `json:"stopReason,omitempty"`
}

// ACPSession 单个用户的 ACP 会话
type ACPSession struct {
	UserId         string
	SessionId      string
	Title          string
	WelcomeMessage string
	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 -> 当前活跃会话
	history    map[string][]SessionListItem // key: userId -> 历史会话列表（元数据）
	mu         sync.RWMutex
	maxIdle    time.Duration
	serviceDir string
	acpConfig  acp.ClientConfig
}

// NewACPManager 创建 ACP 管理器
func NewACPManager(serviceDir string, env []string) *ACPManager {
	mgr := &ACPManager{
		sessions:   make(map[string]*ACPSession),
		history:    make(map[string][]SessionListItem),
		maxIdle:    30 * time.Minute,
		serviceDir: serviceDir,
		acpConfig: acp.ClientConfig{
			Command:    "kiro-cli",
			Args:       []string{"acp", "--agent", "fwh-toolbox-agent"},
			WorkingDir: serviceDir,
			Env:        env,
		},
	}

	// 定时清理空闲会话
	go func() {
		ticker := time.NewTicker(5 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			mgr.CleanupIdle()
		}
	}()

	return mgr
}

// GetOrCreateSession 获取或创建用户的 ACP 会话
func (m *ACPManager) GetOrCreateSession(userId string) (*ACPSession, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	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)
	}

	return m.createNewSessionLocked(userId)
}

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

	// 将当前活跃会话存入历史
	if sess, ok := m.sessions[userId]; ok {
		m.archiveSessionLocked(userId, sess)
		sess.Client.Close()
		delete(m.sessions, userId)
	}

	return m.createNewSessionLocked(userId)
}

// LoadSession 加载历史会话（恢复 AI 上下文）
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.archiveSessionLocked(userId, sess)
		sess.Client.Close()
		delete(m.sessions, userId)
	}

	// 创建新 ACP Client 并加载历史会话
	client, err := acp.NewClient(m.acpConfig)
	if err != nil {
		return nil, err
	}

	if err := client.Initialize(); err != nil {
		client.Close()
		return nil, err
	}

	// 使用 session/load 恢复会话（含回放历史消息）
	loadedId, replayMessages, err := client.LoadSession(sessionId, m.serviceDir)
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("加载会话失败: %w", err)
	}

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

	// 从历史中找到标题
	title := ""
	if userHistory, ok := m.history[userId]; ok {
		for i, item := range userHistory {
			if item.SessionId == sessionId {
				title = item.Title
				// 从历史列表中移除（因为已变为活跃）
				m.history[userId] = append(userHistory[:i], userHistory[i+1:]...)
				break
			}
		}
	}

	sess := &ACPSession{
		UserId:         userId,
		SessionId:      loadedId,
		Title:          title,
		ReplayMessages: replayMessages,
		Client:         client,
		LastActive:     time.Now(),
		CreatedAt:      time.Now(),
	}
	client.UserId = userId

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

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

	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,
		})
	}

	// 历史会话
	if userHistory, ok := m.history[userId]; ok {
		list = append(list, userHistory...)
	}

	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
		return
	}

	// 检查历史会话
	if userHistory, ok := m.history[userId]; ok {
		for i := range userHistory {
			if userHistory[i].SessionId == sessionId {
				userHistory[i].Title = title
				return
			}
		}
	}
}

// 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)
		return
	}

	// 从历史列表中删除
	if userHistory, ok := m.history[userId]; ok {
		for i, item := range userHistory {
			if item.SessionId == sessionId {
				m.history[userId] = append(userHistory[:i], userHistory[i+1:]...)
				return
			}
		}
	}
}

// createNewSessionLocked 内部方法：创建新会话（调用时必须持有 mu 锁）
func (m *ACPManager) createNewSessionLocked(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, 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,
		Client:         client,
		LastActive:     now,
		CreatedAt:      now,
	}
	client.UserId = userId

	m.sessions[userId] = sess
	zap.L().Info("[AIChat] 会话创建成功", zap.String("userId", userId), zap.String("sessionId", sessionId))
	return sess, nil
}

// archiveSessionLocked 内部方法：将会话存入历史（调用时必须持有 mu 锁）
func (m *ACPManager) archiveSessionLocked(userId string, sess *ACPSession) {
	item := SessionListItem{
		SessionId:  sess.SessionId,
		Title:      sess.Title,
		CreatedAt:  sess.CreatedAt.UnixMilli(),
		LastActive: sess.LastActive.UnixMilli(),
		IsActive:   false,
	}

	if _, ok := m.history[userId]; !ok {
		m.history[userId] = []SessionListItem{}
	}

	// 插入到历史列表头部（最近的在前）
	m.history[userId] = append([]SessionListItem{item}, m.history[userId]...)

	// 限制历史数量（最多保留 50 个）
	if len(m.history[userId]) > 50 {
		m.history[userId] = m.history[userId][:50]
	}
}

// Prompt 发送消息并流式读取响应
func (m *ACPManager) Prompt(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
	}

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

	// 转换 acp.Event → AIChatEvent
	for event := range acpEvents {
		switch event.Type {
		case "text":
			eventCh <- AIChatEvent{Type: "text", Content: event.Content}
		case "tool_call":
			eventCh <- AIChatEvent{Type: "tool_call", ToolName: event.ToolName, Status: event.ToolStatus}
		case "tool_update":
			eventCh <- AIChatEvent{Type: "tool_update", ToolName: event.ToolName, Status: event.ToolStatus}
		case "permission":
			eventCh <- AIChatEvent{Type: "permission", Content: toJSON(event.Permission)}
		case "file_created":
			eventCh <- AIChatEvent{Type: "file_created", Content: toJSON(event.File)}
		case "done":
			eventCh <- AIChatEvent{Type: "done", SessionId: event.SessionId, StopReason: event.StopReason}
		case "error":
			eventCh <- AIChatEvent{Type: "error", Content: event.Content}
		}
	}

	close(eventCh)
}

// 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
}

// 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)
		zap.L().Info("[AIChat] 会话已删除", zap.String("userId", userId))
	}
}

// GetSessionInfo 获取用户会话信息
func (m *ACPManager) GetSessionInfo(userId string) *ACPSession {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.sessions[userId]
}

// CleanupIdle 清理空闲超时的会话
func (m *ACPManager) CleanupIdle() {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	for userId, sess := range m.sessions {
		if now.Sub(sess.LastActive) > m.maxIdle {
			sess.Client.Close()
			delete(m.sessions, userId)
			zap.L().Info("[AIChat] 清理空闲会话", zap.String("userId", userId))
		}
	}
}

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