Agent-System/internal/ctxmgr/chat_history.go

package ctxmgr

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"
)

// EnsureChatSession 注册或更新会话元数据（不写入消息）。
func (m *ContextManager) EnsureChatSession(ctx context.Context, sessionID string, userID, projectID int64, title string) (*ChatSessionSummary, error) {
	if sessionID == "" {
		return nil, errors.New("session_id required")
	}
	if title == "" {
		title = "对话协同会话"
	}

	if m.pool != nil {
		var uid *int64
		if userID > 0 {
			uid = &userID
		}
		var pid *int64
		if projectID > 0 {
			pid = &projectID
		}
		_, err := m.pool.Exec(ctx, `
			INSERT INTO chat_sessions (session_id, user_id, project_id, title, created_at, updated_at)
			VALUES ($1, $2, $3, $4, now(), now())
			ON CONFLICT (session_id) DO UPDATE SET
				title = COALESCE(NULLIF($4, ''), chat_sessions.title),
				user_id = COALESCE(chat_sessions.user_id, $2),
				project_id = COALESCE($3, chat_sessions.project_id),
				updated_at = now()
		`, sessionID, uid, pid, title)
		if err != nil {
			return nil, fmt.Errorf("ensure chat session failed: %w", err)
		}

		var summary ChatSessionSummary
		err = m.pool.QueryRow(ctx, `
			SELECT s.session_id, s.title, COALESCE(s.user_id, 0), COALESCE(s.project_id, 0),
			       COALESCE(p.name, ''),
			       COUNT(m.id)::int,
			       COALESCE((SELECT content FROM chat_messages WHERE session_id = s.session_id ORDER BY created_at DESC LIMIT 1), ''),
			       s.updated_at
			FROM chat_sessions s
			LEFT JOIN projects p ON s.project_id = p.id
			LEFT JOIN chat_messages m ON s.session_id = m.session_id
			WHERE s.session_id = $1
			GROUP BY s.session_id, s.title, s.user_id, s.project_id, p.name, s.updated_at
		`, sessionID).Scan(
			&summary.SessionID, &summary.Title, &summary.UserID, &summary.ProjectID,
			&summary.ProjectName, &summary.MessageCount, &summary.LastMessage, &summary.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("load ensured chat session failed: %w", err)
		}
		return &summary, nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	sess, exists := m.sessions[sessionID]
	if !exists {
		sess = &ChatSessionSummary{
			SessionID: sessionID,
			Title:     title,
			UserID:    userID,
			ProjectID: projectID,
			UpdatedAt: time.Now(),
		}
		m.sessions[sessionID] = sess
	} else {
		if title != "" {
			sess.Title = title
		}
		if userID > 0 && sess.UserID <= 0 {
			sess.UserID = userID
		}
		if projectID > 0 {
			sess.ProjectID = projectID
		}
		sess.UpdatedAt = time.Now()
	}
	if msgs, ok := m.messages[sessionID]; ok {
		sess.MessageCount = len(msgs)
		if len(msgs) > 0 {
			sess.LastMessage = msgs[len(msgs)-1].Content
		}
	}
	copy := *sess
	return &copy, nil
}

// ErrSessionOwnerMismatch 表示调用者试图操作已绑定给其他用户的会话。
var ErrSessionOwnerMismatch = errors.New("ctxmgr: session owner mismatch")

// EnsureChatSessionOwned 在写入前校验归属：已绑定且不匹配调用者时返回
// ErrSessionOwnerMismatch；未绑定或本就是调用者本人时正常落地。
func (m *ContextManager) EnsureChatSessionOwned(
	ctx context.Context, sessionID string, callerUserID, projectID int64, title string,
) (*ChatSessionSummary, error) {
	if callerUserID > 0 {
		if owner, err := m.GetSessionOwnerID(ctx, sessionID); err == nil && owner > 0 && owner != callerUserID {
			return nil, ErrSessionOwnerMismatch
		}
	}
	return m.EnsureChatSession(ctx, sessionID, callerUserID, projectID, title)
}

// ListChatSessions 分页查询历史会话列表
func (m *ContextManager) ListChatSessions(ctx context.Context, filter SessionFilter) ([]ChatSessionSummary, error) {
	if filter.Limit <= 0 {
		filter.Limit = 20
	}
	if filter.Offset < 0 {
		filter.Offset = 0
	}

	if m.pool != nil {
		query := `
			SELECT s.session_id, s.title, COALESCE(s.user_id, 0), COALESCE(s.project_id, 0),
			       COALESCE(p.name, ''),
			       COUNT(m.id)::int as msg_count,
			       COALESCE((SELECT content FROM chat_messages WHERE session_id = s.session_id ORDER BY created_at DESC LIMIT 1), '') as last_msg,
			       s.updated_at
			FROM chat_sessions s
			LEFT JOIN projects p ON s.project_id = p.id
			LEFT JOIN chat_messages m ON s.session_id = m.session_id
			WHERE ($1 = 0 OR s.user_id = $1) AND ($2 = 0 OR s.project_id = $2)
			  AND COALESCE(s.status, 'active') <> 'demo_retired'
			  AND COALESCE(p.status, 'active') <> 'demo_retired'
			GROUP BY s.session_id, s.title, s.user_id, s.project_id, p.name, s.updated_at
			ORDER BY s.updated_at DESC
			LIMIT $3 OFFSET $4
		`
		rows, err := m.pool.Query(ctx, query, filter.UserID, filter.ProjectID, filter.Limit, filter.Offset)
		if err != nil {
			return nil, fmt.Errorf("list chat sessions failed: %w", err)
		}
		defer rows.Close()

		var sessions []ChatSessionSummary
		for rows.Next() {
			var item ChatSessionSummary
			if err := rows.Scan(&item.SessionID, &item.Title, &item.UserID, &item.ProjectID, &item.ProjectName, &item.MessageCount, &item.LastMessage, &item.UpdatedAt); err == nil {
				sessions = append(sessions, item)
			}
		}
		return sessions, rows.Err()
	}

	m.mu.RLock()
	defer m.mu.RUnlock()

	var all []ChatSessionSummary
	for _, sess := range m.sessions {
		if filter.UserID > 0 && sess.UserID != filter.UserID {
			continue
		}
		if filter.ProjectID > 0 && sess.ProjectID != filter.ProjectID {
			continue
		}
		status := sess.Status
		if status == "" {
			status = "active"
		}
		if status == "demo_retired" {
			continue
		}
		all = append(all, *sess)
	}

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

	start := filter.Offset
	if start >= len(all) {
		return []ChatSessionSummary{}, nil
	}
	end := start + filter.Limit
	if end > len(all) {
		end = len(all)
	}

	return all[start:end], nil
}

// GetChatHistory 分页拉取指定会话的历史多轮消息记录
func (m *ContextManager) GetChatHistory(ctx context.Context, sessionID string, limit, offset int) ([]ChatMessageRecord, error) {
	if sessionID == "" {
		return nil, errors.New("sessionID required")
	}
	if limit <= 0 {
		limit = 50
	}
	if offset < 0 {
		offset = 0
	}

	if m.pool != nil {
		rows, err := m.pool.Query(ctx, `
			SELECT id, session_id, COALESCE(user_id, 0), role, content, intent, branch, created_at
			FROM chat_messages
			WHERE session_id = $1
			ORDER BY created_at ASC
			LIMIT $2 OFFSET $3
		`, sessionID, limit, offset)
		if err != nil {
			return nil, fmt.Errorf("query chat history failed: %w", err)
		}
		defer rows.Close()

		var msgs []ChatMessageRecord
		for rows.Next() {
			var rec ChatMessageRecord
			if err := rows.Scan(&rec.ID, &rec.SessionID, &rec.UserID, &rec.Role, &rec.Content, &rec.Intent, &rec.Branch, &rec.CreatedAt); err == nil {
				msgs = append(msgs, rec)
			}
		}
		return msgs, rows.Err()
	}

	m.mu.RLock()
	defer m.mu.RUnlock()

	rawMsgs := m.messages[sessionID]
	if offset >= len(rawMsgs) {
		return []ChatMessageRecord{}, nil
	}

	end := offset + limit
	if end > len(rawMsgs) {
		end = len(rawMsgs)
	}

	res := make([]ChatMessageRecord, end-offset)
	copy(res, rawMsgs[offset:end])
	return res, nil
}

// GetSessionProjectID 返回会话关联项目 ID；无项目或未找到会话时返回 0。
func (m *ContextManager) GetSessionProjectID(ctx context.Context, sessionID string) (int64, error) {
	if sessionID == "" {
		return 0, errors.New("session_id required")
	}
	if m.pool != nil {
		var projectID int64
		err := m.pool.QueryRow(ctx, `
			SELECT COALESCE(project_id, 0) FROM chat_sessions WHERE session_id = $1
		`, sessionID).Scan(&projectID)
		if err != nil {
			return 0, fmt.Errorf("session not found: %w", err)
		}
		return projectID, nil
	}

	m.mu.RLock()
	defer m.mu.RUnlock()
	if s, ok := m.sessions[sessionID]; ok {
		return s.ProjectID, nil
	}
	return 0, fmt.Errorf("session not found: %s", sessionID)
}

// GetSessionOwnerID 返回会话归属用户 ID；会话不存在时返回 error。
func (m *ContextManager) GetSessionOwnerID(ctx context.Context, sessionID string) (int64, error) {
	if sessionID == "" {
		return 0, errors.New("session_id required")
	}
	if m.pool != nil {
		var userID int64
		err := m.pool.QueryRow(ctx, `
			SELECT COALESCE(user_id, 0) FROM chat_sessions WHERE session_id = $1
		`, sessionID).Scan(&userID)
		if err != nil {
			return 0, fmt.Errorf("session not found: %w", err)
		}
		if userID <= 0 {
			return 0, errors.New("session owner not set")
		}
		return userID, nil
	}

	m.mu.RLock()
	defer m.mu.RUnlock()
	sess, ok := m.sessions[sessionID]
	if !ok || sess.UserID <= 0 {
		return 0, errors.New("session not found")
	}
	return sess.UserID, nil
}

// RestoreSessionContext 从数据库/内存历史中重构会话上下文
func (m *ContextManager) RestoreSessionContext(ctx context.Context, sessionID string) (*RestoredSessionContext, error) {
	if sessionID == "" {
		return nil, errors.New("sessionID required")
	}

	// 1. 获取全量历史消息
	msgs, err := m.GetChatHistory(ctx, sessionID, 200, 0)
	if err != nil {
		return nil, err
	}

	restored := &RestoredSessionContext{
		SessionID: sessionID,
		Title:     "对话协同会话",
		Messages:  msgs,
		Facts:     make(map[string]string),
	}

	// 2. 从消息或会话表中提取归属用户与项目
	if m.pool != nil {
		var userID, projectID int64
		var title string
		err := m.pool.QueryRow(ctx, `
			SELECT COALESCE(user_id, 0), COALESCE(project_id, 0), title
			FROM chat_sessions WHERE session_id = $1
		`, sessionID).Scan(&userID, &projectID, &title)
		if err == nil {
			restored.UserID = userID
			restored.ProjectID = projectID
			if title != "" {
				restored.Title = title
			}
		}
	} else {
		m.mu.RLock()
		if s, ok := m.sessions[sessionID]; ok {
			restored.UserID = s.UserID
			restored.ProjectID = s.ProjectID
			restored.Title = s.Title
		}
		m.mu.RUnlock()
	}

	// 3. 若有项目ID，同步抽取黑板已有事实
	if restored.ProjectID > 0 {
		pCtx, err := m.GetProjectContext(ctx, restored.ProjectID)
		if err == nil && pCtx != nil {
			for k, v := range pCtx.BlackboardFacts {
				restored.Facts[k] = v
			}
		}
	}

	// 4. 遍历历史消息尝试提炼显式槽位或事实
	for _, msg := range msgs {
		if msg.UserID > 0 && restored.UserID == 0 {
			restored.UserID = msg.UserID
		}
		// 简单提取器械与市场关键词（若未固化）
		if strings.Contains(msg.Content, "超声") && restored.Facts["device_type"] == "" {
			restored.Facts["device_type"] = "超声手术设备"
		}
		if strings.Contains(msg.Content, "内窥镜") && restored.Facts["device_type"] == "" {
			restored.Facts["device_type"] = "内窥镜系统"
		}
	}

	return restored, nil
}