CamTalk/backend/internal/session/memory.go

package session

import (
	"context"
	"encoding/json"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"

	"github.com/hhs/camtalk/internal/logger"
	"github.com/hhs/camtalk/internal/models"
	"github.com/hhs/camtalk/internal/store"
)

const (
	defaultTTL         = 30 * time.Minute
	defaultHistorySize = 20
)

// sessionEntry 内部会话条目。
type sessionEntry struct {
	session     models.Session
	history     []models.Message
	activeReqID string
	lastActive  time.Time
}

// MemoryManager 基于内存的 SessionManager 实现。
// 适用于 MVP 和无 Redis 的开发环境。
type MemoryManager struct {
	mu          sync.RWMutex
	sessions    map[string]*sessionEntry
	ttl         time.Duration
	maxHistory  int
	stopCleaner chan struct{}
	msgRepo     store.MessageRepository  // 可选，消息持久化（Write-Through）
	sessRepo    store.SessionRepository  // 可选，会话持久化（Write-Through）
}

// Option MemoryManager 的函数式选项。
type Option func(*MemoryManager)

// WithMessageRepository 注入消息持久化仓库，启用 Write-Through 模式。
func WithMessageRepository(repo store.MessageRepository) Option {
	return func(m *MemoryManager) {
		m.msgRepo = repo
	}
}

// WithSessionRepository 注入会话持久化仓库，启用会话元数据 Write-Through 模式。
func WithSessionRepository(repo store.SessionRepository) Option {
	return func(m *MemoryManager) {
		m.sessRepo = repo
	}
}

// NewMemoryManager 创建内存版 SessionManager。
// ttl 为会话过期时间，maxHistory 为对话历史上限（0 表示使用默认值 20）。
// opts 为可选配置，如 WithMessageRepository 启用消息持久化。
func NewMemoryManager(ttl time.Duration, maxHistory int, opts ...Option) *MemoryManager {
	if ttl <= 0 {
		ttl = defaultTTL
	}
	if maxHistory <= 0 {
		maxHistory = defaultHistorySize
	}

	m := &MemoryManager{
		sessions:    make(map[string]*sessionEntry),
		ttl:         ttl,
		maxHistory:  maxHistory,
		stopCleaner: make(chan struct{}),
	}

	for _, opt := range opts {
		opt(m)
	}

	// 启动后台清理 goroutine，每分钟清除过期会话。
	go m.cleanLoop()

	return m
}

// cleanLoop 后台定期清理过期会话。
func (m *MemoryManager) cleanLoop() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			m.cleanExpired()
		case <-m.stopCleaner:
			return
		}
	}
}

// cleanExpired 清除所有过期会话。
func (m *MemoryManager) cleanExpired() {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	for id, entry := range m.sessions {
		if now.Sub(entry.lastActive) > m.ttl {
			delete(m.sessions, id)
			logger.Log.Debugw("session expired (cleaner)", "session", id)
		}
	}
}

// Stop 停止后台清理 goroutine。应用退出前调用。
func (m *MemoryManager) Stop() {
	close(m.stopCleaner)
}

// isExpired 检查会话是否过期（调用方需持锁或在已知 entry 存在时调用）。
func (m *MemoryManager) isExpired(entry *sessionEntry) bool {
	return time.Since(entry.lastActive) > m.ttl
}

// Create 创建新会话。userID 为空表示匿名会话。
func (m *MemoryManager) Create(ctx context.Context, userID string, config models.SessionConfig) (string, error) {
	m.mu.Lock()

	id := uuid.New().String()
	now := time.Now()
	m.sessions[id] = &sessionEntry{
		session: models.Session{
			ID:        id,
			UserID:    userID,
			Title:     models.DefaultSessionTitle,
			CreatedAt: now,
			UpdatedAt: now,
			Config:    config,
		},
		history:    make([]models.Message, 0),
		lastActive: now,
	}
	m.mu.Unlock()

	// Write-Through：异步写 PG
	if m.sessRepo != nil {
		go func() {
			cfgJSON, _ := json.Marshal(config)
			if err := m.sessRepo.Save(ctx, store.SessionRecord{
				ID: id, UserID: userID, Title: models.DefaultSessionTitle,
				Config: cfgJSON, CreatedAt: now, UpdatedAt: now,
			}); err != nil {
				logger.Log.Warnw("persist session failed", "session", id, "error", err)
			}
		}()
	}

	logger.Log.Debugw("session created", "session", id, "user_id", userID)
	return id, nil
}

// Get 获取会话。内存中不存在时，尝试从 PG 加载（透明恢复）。
func (m *MemoryManager) Get(ctx context.Context, sessionID string) (*models.Session, error) {
	m.mu.RLock()
	entry, ok := m.sessions[sessionID]
	if ok && !m.isExpired(entry) {
		sess := entry.session
		m.mu.RUnlock()
		return &sess, nil
	}
	m.mu.RUnlock()

	// 内存未命中，尝试从 PG 加载
	if m.sessRepo != nil {
		rec, err := m.sessRepo.FindByID(ctx, sessionID)
		if err != nil {
			return nil, ErrSessionNotFound
		}
		sess := m.recordToSession(rec)
		// 加载到内存（含消息历史）
		if m.msgRepo != nil {
			_ = m.LoadSessionFromRepo(ctx, sess)
		} else {
			_ = m.LoadSession(sess, nil)
		}
		return sess, nil
	}

	return nil, ErrSessionNotFound
}

// UpdateConfig 更新会话配置。
func (m *MemoryManager) UpdateConfig(ctx context.Context, sessionID string, patch models.SessionConfigPatch) error {
	m.mu.Lock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		m.mu.Unlock()
		return ErrSessionNotFound
	}

	patch.Apply(&entry.session.Config)
	entry.lastActive = time.Now()
	cfg := entry.session.Config
	m.mu.Unlock()

	// Write-Through：异步更新 PG
	if m.sessRepo != nil {
		go func() {
			cfgJSON, _ := json.Marshal(cfg)
			if err := m.sessRepo.UpdateConfig(ctx, sessionID, cfgJSON); err != nil {
				logger.Log.Warnw("update session config in DB failed", "session", sessionID, "error", err)
			}
		}()
	}

	logger.Log.Debugw("session config updated", "session", sessionID)
	return nil
}

// UpdateTitle 更新会话标题。
func (m *MemoryManager) UpdateTitle(ctx context.Context, sessionID string, title string) error {
	m.mu.Lock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		m.mu.Unlock()
		return ErrSessionNotFound
	}

	entry.session.Title = title
	entry.session.UpdatedAt = time.Now()
	entry.lastActive = time.Now()
	m.mu.Unlock()

	// Write-Through：异步更新 PG
	if m.sessRepo != nil {
		go func() {
			if err := m.sessRepo.UpdateTitle(ctx, sessionID, title); err != nil {
				logger.Log.Warnw("update session title in DB failed", "session", sessionID, "error", err)
			}
		}()
	}

	logger.Log.Debugw("session title updated", "session", sessionID, "title", title)
	return nil
}

// ListByUser 获取用户的对话列表（分页，按 UpdatedAt 降序）。
// 若配置了 SessionRepository，从 PG 查询（包含内存中已过期的会话）。
// 若配置了 MessageRepository，消息统计从 PostgreSQL 聚合查询（更准确）。
func (m *MemoryManager) ListByUser(ctx context.Context, userID string, page, size int) ([]ConversationSummary, int, error) {
	if page <= 0 {
		page = 1
	}
	if size <= 0 {
		size = 20
	}

	// 优先从 PG 查询会话列表（包含已过期的会话）
	if m.sessRepo != nil {
		recs, total, err := m.sessRepo.FindByUser(ctx, userID, page, size)
		if err != nil {
			logger.Log.Warnw("list sessions from DB failed, falling back to in-memory", "error", err)
			return m.listByUserFromMemory(ctx, userID, page, size)
		}

		list := make([]ConversationSummary, 0, len(recs))
		var sessionIDs []string
		for _, rec := range recs {
			list = append(list, ConversationSummary{
				ID:        rec.ID,
				Title:     rec.Title,
				UpdatedAt: rec.UpdatedAt,
			})
			sessionIDs = append(sessionIDs, rec.ID)
		}

		// 用内存中的消息数填充
		m.mu.RLock()
		for i := range list {
			if entry, ok := m.sessions[list[i].ID]; ok {
				list[i].MessageCount = len(entry.history)
				if len(entry.history) > 0 {
					list[i].LastMessage = entry.history[len(entry.history)-1].Content
				}
			}
		}
		m.mu.RUnlock()

		// 从 PG 获取更准确的消息统计
		if m.msgRepo != nil && len(sessionIDs) > 0 {
			if stats, err := m.msgRepo.GetSessionMessageStats(ctx, sessionIDs); err == nil {
				for i := range list {
					if s, ok := stats[list[i].ID]; ok {
						list[i].LastMessage = s.LastMessage
						list[i].MessageCount = s.MessageCount
					}
				}
			}
		}

		return list, total, nil
	}

	// fallback：纯内存查询
	return m.listByUserFromMemory(ctx, userID, page, size)
}

// listByUserFromMemory 从内存中获取用户的对话列表（无 PG 时的 fallback）。
func (m *MemoryManager) listByUserFromMemory(ctx context.Context, userID string, page, size int) ([]ConversationSummary, int, error) {
	m.mu.RLock()

	var list []ConversationSummary
	var sessionIDs []string
	for _, entry := range m.sessions {
		if entry.session.UserID != userID || m.isExpired(entry) {
			continue
		}
		summary := ConversationSummary{
			ID:        entry.session.ID,
			Title:     entry.session.Title,
			UpdatedAt: entry.lastActive,
		}
		summary.MessageCount = len(entry.history)
		if len(entry.history) > 0 {
			summary.LastMessage = entry.history[len(entry.history)-1].Content
		}
		list = append(list, summary)
		sessionIDs = append(sessionIDs, entry.session.ID)
	}
	m.mu.RUnlock()

	// 从 PG 获取更准确的消息统计
	if m.msgRepo != nil && len(sessionIDs) > 0 {
		if stats, err := m.msgRepo.GetSessionMessageStats(ctx, sessionIDs); err == nil {
			for i := range list {
				if s, ok := stats[list[i].ID]; ok {
					list[i].LastMessage = s.LastMessage
					list[i].MessageCount = s.MessageCount
				}
			}
		}
	}

	sort.Slice(list, func(i, j int) bool {
		return list[i].UpdatedAt.After(list[j].UpdatedAt)
	})

	total := len(list)
	start := (page - 1) * size
	if start >= total {
		return []ConversationSummary{}, total, nil
	}
	end := start + size
	if end > total {
		end = total
	}

	return list[start:end], total, nil
}

// GetHistory 获取最近 N 轮对话历史。
func (m *MemoryManager) GetHistory(_ context.Context, sessionID string, limit int) ([]models.Message, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		return nil, ErrSessionNotFound
	}

	if limit <= 0 || limit > len(entry.history) {
		limit = len(entry.history)
	}

	// 返回最近 limit 条的副本
	result := make([]models.Message, limit)
	copy(result, entry.history[len(entry.history)-limit:])
	return result, nil
}

// AppendMessage 追加一条对话消息，同时刷新 TTL。
// 若配置了 MessageRepository，消息会异步写入 PostgreSQL（Write-Through）。
func (m *MemoryManager) AppendMessage(_ context.Context, sessionID string, msg models.Message) error {
	m.mu.Lock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		m.mu.Unlock()
		return ErrSessionNotFound
	}

	entry.history = append(entry.history, msg)

	// 自动更新标题：首条 user 消息时，如果标题为默认值，自动更新为消息前 20 字符
	if msg.Role == "user" && entry.session.Title == models.DefaultSessionTitle {
		entry.session.Title = generateTitle(msg.Content)
	}

	// 超过上限时裁剪，保留最新的 maxHistory 条
	if len(entry.history) > m.maxHistory {
		entry.history = entry.history[len(entry.history)-m.maxHistory:]
	}

	now := time.Now()
	entry.lastActive = now
	entry.session.UpdatedAt = now
	m.mu.Unlock()

	// Write-Through：异步写冷存储，不阻塞调用方
	if m.msgRepo != nil {
		go func() {
			if err := m.msgRepo.SaveMessage(context.Background(), sessionID, msg, 0); err != nil {
				logger.Log.Warnw("persist message failed", "session", sessionID, "error", err)
			}
		}()
	}

	return nil
}

// generateTitle 从首条消息生成对话标题（取前 20 个字符）。
func generateTitle(firstMessage string) string {
	runes := []rune(firstMessage)
	if len(runes) > 20 {
		return string(runes[:20]) + "…"
	}
	return firstMessage
}

// LoadSession 从外部存储加载会话到内存热存储。
// 用于 conversation_id 恢复场景：WS 连接时会话不在内存中，从 PostgreSQL 加载。
// 若会话已在内存中，返回 nil（幂等）。
func (m *MemoryManager) LoadSession(sess *models.Session, messages []models.Message) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.sessions[sess.ID]; ok {
		return nil // 已在内存中，无需重复加载
	}

	m.sessions[sess.ID] = &sessionEntry{
		session:    *sess,
		history:    messages,
		lastActive: time.Now(),
	}

	logger.Log.Debugw("session loaded from DB", "session", sess.ID, "messages", len(messages))
	return nil
}

// LoadSessionFromRepo 从 MessageRepository 加载会话消息并注册到内存。
// 适用于已注入 MessageRepository 的场景，调用方只需传入 session 元数据。
func (m *MemoryManager) LoadSessionFromRepo(ctx context.Context, sess *models.Session) error {
	if m.msgRepo == nil {
		return m.LoadSession(sess, nil)
	}

	// 从冷存储加载全部消息（limit=0 表示全量）
	stored, err := m.msgRepo.GetMessages(ctx, sess.ID, 0, 0)
	if err != nil {
		return err
	}

	messages := make([]models.Message, len(stored))
	for i, s := range stored {
		messages[i] = models.Message{Role: s.Role, Content: s.Content}
	}

	return m.LoadSession(sess, messages)
}

// SetActiveRequest 标记当前正在处理的请求 ID。
func (m *MemoryManager) SetActiveRequest(_ context.Context, sessionID string, requestID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		return ErrSessionNotFound
	}

	entry.activeReqID = requestID
	entry.lastActive = time.Now()
	return nil
}

// GetActiveRequestID 获取当前活跃请求 ID。
func (m *MemoryManager) GetActiveRequestID(_ context.Context, sessionID string) (string, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		return "", ErrSessionNotFound
	}

	return entry.activeReqID, nil
}

// ClearActiveRequest 清除活跃请求标记。
func (m *MemoryManager) ClearActiveRequest(_ context.Context, sessionID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		return ErrSessionNotFound
	}

	entry.activeReqID = ""
	entry.lastActive = time.Now()
	return nil
}

// Touch 刷新 TTL。
func (m *MemoryManager) Touch(_ context.Context, sessionID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	entry, ok := m.sessions[sessionID]
	if !ok || m.isExpired(entry) {
		return ErrSessionNotFound
	}

	entry.lastActive = time.Now()
	return nil
}

// Destroy 显式销毁会话。
func (m *MemoryManager) Destroy(ctx context.Context, sessionID string) error {
	m.mu.Lock()

	if _, ok := m.sessions[sessionID]; !ok {
		m.mu.Unlock()
		return ErrSessionNotFound
	}

	delete(m.sessions, sessionID)
	m.mu.Unlock()

	// Write-Through：异步删除 PG
	if m.sessRepo != nil {
		go func() {
			if err := m.sessRepo.Delete(ctx, sessionID); err != nil {
				logger.Log.Warnw("delete session from DB failed", "session", sessionID, "error", err)
			}
		}()
	}

	logger.Log.Debugw("session destroyed", "session", sessionID)
	return nil
}

// ActiveCount 返回当前活跃会话数。
func (m *MemoryManager) ActiveCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()

	now := time.Now()
	count := 0
	for _, entry := range m.sessions {
		if now.Sub(entry.lastActive) <= m.ttl {
			count++
		}
	}
	return count
}

// recordToSession 将 store.SessionRecord 转换为 models.Session。
func (m *MemoryManager) recordToSession(rec *store.SessionRecord) *models.Session {
	cfg := models.DefaultConfig()
	if len(rec.Config) > 0 {
		_ = json.Unmarshal(rec.Config, &cfg)
	}
	return &models.Session{
		ID:        rec.ID,
		UserID:    rec.UserID,
		Title:     rec.Title,
		CreatedAt: rec.CreatedAt,
		UpdatedAt: rec.UpdatedAt,
		Config:    cfg,
	}
}