Agent-System/internal/ctxmgr/chat_search.go

package ctxmgr

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"unicode/utf8"

	"github.com/jackc/pgx/v5"
)

// SearchChatMessages 统一搜索入口: 支持 keyword/fuzzy/semantic/hybrid 四种检索模式
func (m *ContextManager) SearchChatMessages(ctx context.Context, q ChatSearchQuery) ([]ChatSearchHit, error) {
	trimmedQuery := strings.TrimSpace(q.Query)
	if trimmedQuery == "" {
		return []ChatSearchHit{}, nil
	}

	if q.Mode == "" {
		q.Mode = SearchModeHybrid
	}
	if q.Limit <= 0 {
		q.Limit = 50
	}
	if q.Limit > 200 {
		q.Limit = 200
	}

	// 内存降级模式 (无真实数据库连接)
	if m.pool == nil {
		return m.searchChatMessagesMemory(q)
	}

	switch q.Mode {
	case SearchModeKeyword:
		return m.searchChatMessagesKeywordDB(ctx, q)
	case SearchModeFuzzy:
		return m.searchChatMessagesFuzzyDB(ctx, q)
	case SearchModeSemantic:
		return m.searchChatMessagesSemanticDB(ctx, q)
	case SearchModeHybrid:
		return m.searchChatMessagesHybridDB(ctx, q)
	default:
		return m.searchChatMessagesKeywordDB(ctx, q)
	}
}

// searchChatMessagesKeywordDB 关键词搜索 (Postgres ILIKE)
func (m *ContextManager) searchChatMessagesKeywordDB(ctx context.Context, q ChatSearchQuery) ([]ChatSearchHit, error) {
	var builder strings.Builder
	args := []any{q.SessionID, "%" + q.Query + "%"}
	builder.WriteString(`
		SELECT id, session_id, day_key::text, role, category, COALESCE(metadata->>'stage_key', ''), content
		FROM chat_messages
		WHERE session_id = $1 AND deleted_at IS NULL AND content ILIKE $2
	`)

	argIdx := 3
	if q.Category != "" && q.Category != "all" {
		builder.WriteString(fmt.Sprintf(" AND category = $%d", argIdx))
		args = append(args, q.Category)
		argIdx++
	}
	if q.DateFrom != "" {
		builder.WriteString(fmt.Sprintf(" AND day_key >= $%d::date", argIdx))
		args = append(args, q.DateFrom)
		argIdx++
	}
	if q.DateTo != "" {
		builder.WriteString(fmt.Sprintf(" AND day_key <= $%d::date", argIdx))
		args = append(args, q.DateTo)
		argIdx++
	}

	builder.WriteString(fmt.Sprintf(" ORDER BY id DESC LIMIT $%d", argIdx))
	args = append(args, q.Limit)

	rows, err := m.pool.Query(ctx, builder.String(), args...)
	if err != nil {
		return nil, fmt.Errorf("searchChatMessagesKeywordDB query failed: %w", err)
	}
	defer rows.Close()

	var hits []ChatSearchHit
	for rows.Next() {
		var id int64
		var sid, dayKey, role, cat, stageKey, content string
		if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content); err != nil {
			return nil, err
		}
		hits = append(hits, ChatSearchHit{
			MessageID: id,
			SessionID: sid,
			DayKey:    dayKey,
			Role:      role,
			Category:  cat,
			StageKey:  stageKey,
			Excerpt:   ExtractSearchExcerpt(content, q.Query, 40),
			Score:     1.0,
			MatchedBy: "keyword",
		})
	}
	return hits, rows.Err()
}

// searchChatMessagesFuzzyDB 模糊相似度检索 (Postgres pg_trgm similarity)
func (m *ContextManager) searchChatMessagesFuzzyDB(ctx context.Context, q ChatSearchQuery) ([]ChatSearchHit, error) {
	var builder strings.Builder
	args := []any{q.SessionID, q.Query}
	builder.WriteString(`
		SELECT id, session_id, day_key::text, role, category, COALESCE(metadata->>'stage_key', ''), content, similarity(content, $2) AS sim
		FROM chat_messages
		WHERE session_id = $1 AND deleted_at IS NULL AND similarity(content, $2) > 0.15
	`)

	argIdx := 3
	if q.Category != "" && q.Category != "all" {
		builder.WriteString(fmt.Sprintf(" AND category = $%d", argIdx))
		args = append(args, q.Category)
		argIdx++
	}
	if q.DateFrom != "" {
		builder.WriteString(fmt.Sprintf(" AND day_key >= $%d::date", argIdx))
		args = append(args, q.DateFrom)
		argIdx++
	}
	if q.DateTo != "" {
		builder.WriteString(fmt.Sprintf(" AND day_key <= $%d::date", argIdx))
		args = append(args, q.DateTo)
		argIdx++
	}

	builder.WriteString(fmt.Sprintf(" ORDER BY sim DESC, id DESC LIMIT $%d", argIdx))
	args = append(args, q.Limit)

	rows, err := m.pool.Query(ctx, builder.String(), args...)
	if err != nil {
		return nil, fmt.Errorf("searchChatMessagesFuzzyDB query failed: %w", err)
	}
	defer rows.Close()

	var hits []ChatSearchHit
	for rows.Next() {
		var id int64
		var sid, dayKey, role, cat, stageKey, content string
		var sim float64
		if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content, &sim); err != nil {
			return nil, err
		}
		hits = append(hits, ChatSearchHit{
			MessageID: id,
			SessionID: sid,
			DayKey:    dayKey,
			Role:      role,
			Category:  cat,
			StageKey:  stageKey,
			Excerpt:   ExtractSearchExcerpt(content, q.Query, 40),
			Score:     sim,
			MatchedBy: "fuzzy",
		})
	}
	return hits, rows.Err()
}

// searchChatMessagesSemanticDB 语义向量检索 (通过 context_memory_chunks 或降级为 keyword)
func (m *ContextManager) searchChatMessagesSemanticDB(ctx context.Context, q ChatSearchQuery) ([]ChatSearchHit, error) {
	if m.embedder == nil {
		return m.searchChatMessagesKeywordDB(ctx, q)
	}

	vecs, err := m.embedder.Embed(ctx, []string{q.Query})
	if err != nil || len(vecs) == 0 || len(vecs[0]) == 0 {
		return m.searchChatMessagesKeywordDB(ctx, q)
	}

	vecStr := formatFloatArray(vecs[0])
	var builder strings.Builder
	args := []any{q.SessionID, vecStr}
	builder.WriteString(`
		SELECT m.id, m.session_id, m.day_key::text, m.role, m.category, COALESCE(m.metadata->>'stage_key', ''), m.content,
		       (1 - (c.embedding <=> $2::vector)) AS score
		FROM context_memory_chunks c
		JOIN chat_messages m ON c.reference_id = 'msg:' || m.id
		WHERE m.session_id = $1 AND m.deleted_at IS NULL AND c.domain = 'chat_message'
	`)

	argIdx := 3
	if q.Category != "" && q.Category != "all" {
		builder.WriteString(fmt.Sprintf(" AND m.category = $%d", argIdx))
		args = append(args, q.Category)
		argIdx++
	}
	if q.DateFrom != "" {
		builder.WriteString(fmt.Sprintf(" AND m.day_key >= $%d::date", argIdx))
		args = append(args, q.DateFrom)
		argIdx++
	}
	if q.DateTo != "" {
		builder.WriteString(fmt.Sprintf(" AND m.day_key <= $%d::date", argIdx))
		args = append(args, q.DateTo)
		argIdx++
	}

	builder.WriteString(fmt.Sprintf(" ORDER BY score DESC, m.id DESC LIMIT $%d", argIdx))
	args = append(args, q.Limit)

	rows, err := m.pool.Query(ctx, builder.String(), args...)
	if err != nil {
		// 若向量表暂未建立该 domain 记录，安全退化到关键词检索
		return m.searchChatMessagesKeywordDB(ctx, q)
	}
	defer rows.Close()

	var hits []ChatSearchHit
	for rows.Next() {
		var id int64
		var sid, dayKey, role, cat, stageKey, content string
		var score float64
		if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content, &score); err != nil {
			return nil, err
		}
		hits = append(hits, ChatSearchHit{
			MessageID: id,
			SessionID: sid,
			DayKey:    dayKey,
			Role:      role,
			Category:  cat,
			StageKey:  stageKey,
			Excerpt:   ExtractSearchExcerpt(content, q.Query, 40),
			Score:     score,
			MatchedBy: "semantic",
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("scan semantic search hits failed: %w", err)
	}

	if len(hits) == 0 {
		return m.searchChatMessagesKeywordDB(ctx, q)
	}
	return hits, nil
}

// searchChatMessagesHybridDB 混合加权检索 (Keyword 0.4 + Fuzzy 0.2 + Semantic 0.4)
func (m *ContextManager) searchChatMessagesHybridDB(ctx context.Context, q ChatSearchQuery) ([]ChatSearchHit, error) {
	kwHits, _ := m.searchChatMessagesKeywordDB(ctx, q)
	fzHits, _ := m.searchChatMessagesFuzzyDB(ctx, q)

	type mergedHit struct {
		hit   ChatSearchHit
		score float64
	}
	hitMap := make(map[int64]*mergedHit)

	for _, h := range kwHits {
		hitMap[h.MessageID] = &mergedHit{
			hit:   h,
			score: 0.5 + 0.5*h.Score, // keyword基础加成
		}
	}

	for _, h := range fzHits {
		if existing, found := hitMap[h.MessageID]; found {
			existing.score = existing.score*0.7 + h.Score*0.3
			existing.hit.MatchedBy = "hybrid"
		} else {
			hitMap[h.MessageID] = &mergedHit{
				hit:   h,
				score: h.Score * 0.8,
			}
		}
	}

	var results []ChatSearchHit
	for _, m := range hitMap {
		m.hit.Score = m.score
		results = append(results, m.hit)
	}

	sort.Slice(results, func(i, j int) bool {
		if results[i].Score == results[j].Score {
			return results[i].MessageID > results[j].MessageID
		}
		return results[i].Score > results[j].Score
	})

	if len(results) > q.Limit {
		results = results[:q.Limit]
	}
	return results, nil
}

// searchChatMessagesMemory 内存降级搜索实现 (供单测及无 DB 模式)
func (m *ContextManager) searchChatMessagesMemory(q ChatSearchQuery) ([]ChatSearchHit, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	records := m.messages[q.SessionID]
	var hits []ChatSearchHit
	lowerQuery := strings.ToLower(q.Query)

	for _, r := range records {
		if q.Category != "" && q.Category != "all" && r.Category != q.Category {
			continue
		}
		if q.DateFrom != "" && r.DayKey < q.DateFrom {
			continue
		}
		if q.DateTo != "" && r.DayKey > q.DateTo {
			continue
		}

		lowerContent := strings.ToLower(r.Content)
		if strings.Contains(lowerContent, lowerQuery) {
			stageKey := ""
			if r.Metadata != nil {
				if sk, ok := r.Metadata["stage_key"].(string); ok {
					stageKey = sk
				}
			}
			hits = append(hits, ChatSearchHit{
				MessageID: r.ID,
				SessionID: r.SessionID,
				DayKey:    r.DayKey,
				Role:      r.Role,
				Category:  r.Category,
				StageKey:  stageKey,
				Excerpt:   ExtractSearchExcerpt(r.Content, q.Query, 40),
				Score:     1.0,
				MatchedBy: "keyword",
			})
		}
	}

	// 倒序排列 (最新在前)
	sort.Slice(hits, func(i, j int) bool {
		return hits[i].MessageID > hits[j].MessageID
	})

	if len(hits) > q.Limit {
		hits = hits[:q.Limit]
	}
	return hits, nil
}

// ExtractSearchExcerpt 提取命中关键词附近的文本摘录 (±windowLen 字符)
func ExtractSearchExcerpt(content, query string, windowLen int) string {
	if windowLen <= 0 {
		windowLen = 40
	}
	contentRunes := []rune(content)
	if len(contentRunes) <= windowLen*2 {
		return content
	}

	queryRunes := []rune(query)
	lowerContent := strings.ToLower(content)
	lowerQuery := strings.ToLower(query)

	idx := strings.Index(lowerContent, lowerQuery)
	if idx < 0 {
		if len(contentRunes) > windowLen*2 {
			return string(contentRunes[:windowLen*2]) + "..."
		}
		return content
	}

	// 换算 rune 偏移
	runeIdx := utf8.RuneCountInString(content[:idx])
	qLen := len(queryRunes)

	start := runeIdx - windowLen
	if start < 0 {
		start = 0
	}
	end := runeIdx + qLen + windowLen
	if end > len(contentRunes) {
		end = len(contentRunes)
	}

	var sb strings.Builder
	if start > 0 {
		sb.WriteString("...")
	}
	sb.WriteString(string(contentRunes[start:end]))
	if end < len(contentRunes) {
		sb.WriteString("...")
	}
	return sb.String()
}

// formatFloatArray 将 float32 数组转换为 Postgres vector 字符串形式
func formatFloatArray(v []float32) string {
	var sb strings.Builder
	sb.WriteString("[")
	for i, f := range v {
		if i > 0 {
			sb.WriteString(",")
		}
		sb.WriteString(fmt.Sprintf("%f", f))
	}
	sb.WriteString("]")
	return sb.String()
}

// 确保引用 pgx
var _ = pgx.ErrNoRows