LaodingBot/internal/memory/store_sqlite.go

package memory

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"laodingbot/internal/logger"

	_ "modernc.org/sqlite"
)

type Message struct {
	ID        int64
	ChatID    string
	UserID    string
	Role      string
	Content   string
	CreatedAt time.Time
}

type SQLiteStore struct {
	db  *sql.DB
	log *logger.Logger
}

func NewSQLiteStore(path string, log *logger.Logger) (*SQLiteStore, error) {
	abs, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}
	if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", abs)
	if err != nil {
		return nil, err
	}
	store := &SQLiteStore{db: db, log: log}
	if err := store.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	if log != nil {
		log.Infof("sqlite store initialized path=%s", abs)
	}
	return store, nil
}

func (s *SQLiteStore) Close() error {
	return s.db.Close()
}

func (s *SQLiteStore) SaveMessage(chatID, userID, role, content string) error {
	if s.log != nil {
		s.log.Debugf("save message chat_id=%s role=%s content_len=%d", chatID, role, len(content))
	}
	_, err := s.db.Exec(`
		INSERT INTO messages(chat_id, user_id, role, content, created_at)
		VALUES (?, ?, ?, ?, ?)
	`, chatID, userID, role, content, time.Now().UTC())
	if err != nil && s.log != nil {
		s.log.Errorf("save message failed chat_id=%s role=%s err=%v", chatID, role, err)
	}
	return err
}

func (s *SQLiteStore) LoadRecent(chatID string, limit int) ([]Message, error) {
	if limit <= 0 {
		limit = 20
	}
	rows, err := s.db.Query(`
		SELECT id, chat_id, user_id, role, content, created_at
		FROM messages
		WHERE chat_id = ?
		ORDER BY id DESC
		LIMIT ?
	`, chatID, limit)
	if err != nil {
		if s.log != nil {
			s.log.Errorf("load recent query failed chat_id=%s err=%v", chatID, err)
		}
		return nil, err
	}
	defer rows.Close()

	messages := make([]Message, 0, limit)
	for rows.Next() {
		var m Message
		if err := rows.Scan(&m.ID, &m.ChatID, &m.UserID, &m.Role, &m.Content, &m.CreatedAt); err != nil {
			return nil, err
		}
		messages = append(messages, m)
	}
	if err := rows.Err(); err != nil {
		if s.log != nil {
			s.log.Errorf("load recent row iteration failed chat_id=%s err=%v", chatID, err)
		}
		return nil, err
	}

	for left, right := 0, len(messages)-1; left < right; left, right = left+1, right-1 {
		messages[left], messages[right] = messages[right], messages[left]
	}
	if s.log != nil {
		s.log.Debugf("load recent success chat_id=%s count=%d", chatID, len(messages))
	}
	return messages, nil
}

func (s *SQLiteStore) SaveCapabilityGap(chatID, userID, intent, reason string) error {
	_, err := s.db.Exec(`
		INSERT INTO capability_gaps(chat_id, user_id, intent, reason, created_at)
		VALUES (?, ?, ?, ?, ?)
	`, chatID, userID, intent, reason, time.Now().UTC())
	if err != nil && s.log != nil {
		s.log.Errorf("save capability gap failed chat_id=%s user_id=%s err=%v", chatID, userID, err)
	}
	return err
}

func (s *SQLiteStore) TopCapabilityGaps(limit int) ([]CapabilityGap, error) {
	if limit <= 0 {
		limit = 20
	}
	rows, err := s.db.Query(`
		SELECT id, chat_id, user_id, intent, reason, created_at
		FROM capability_gaps
		ORDER BY id DESC
		LIMIT ?
	`, limit)
	if err != nil {
		if s.log != nil {
			s.log.Errorf("top capability gaps query failed err=%v", err)
		}
		return nil, err
	}
	defer rows.Close()

	out := make([]CapabilityGap, 0, limit)
	for rows.Next() {
		var item CapabilityGap
		if err := rows.Scan(&item.ID, &item.ChatID, &item.UserID, &item.Intent, &item.Reason, &item.CreatedAt); err != nil {
			return nil, err
		}
		out = append(out, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return out, nil
}

func (s *SQLiteStore) TopCapabilityGapClusters(limit int, since time.Time) ([]CapabilityGapCluster, error) {
	if limit <= 0 {
		limit = 20
	}
	if since.IsZero() {
		since = time.Now().UTC().Add(-7 * 24 * time.Hour)
	}

	rows, err := s.db.Query(`
		SELECT intent, reason, created_at
		FROM capability_gaps
		WHERE created_at >= ?
	`, since)
	if err != nil {
		if s.log != nil {
			s.log.Errorf("top capability gap clusters query failed err=%v", err)
		}
		return nil, err
	}
	defer rows.Close()

	type groupKey struct {
		IntentKey string
		Reason    string
	}
	type agg struct {
		cluster CapabilityGapCluster
	}

	groups := map[groupKey]agg{}
	for rows.Next() {
		var intent string
		var reason string
		var createdAt time.Time
		if err := rows.Scan(&intent, &reason, &createdAt); err != nil {
			return nil, err
		}
		intentKey := normalizeIntentKey(intent)
		reason = strings.TrimSpace(reason)
		k := groupKey{IntentKey: intentKey, Reason: reason}
		current, ok := groups[k]
		if !ok {
			current = agg{cluster: CapabilityGapCluster{
				IntentKey:    intentKey,
				SampleIntent: strings.TrimSpace(intent),
				Reason:       reason,
				Count:        0,
				LastSeenAt:   createdAt,
			}}
		}
		current.cluster.Count++
		if createdAt.After(current.cluster.LastSeenAt) {
			current.cluster.LastSeenAt = createdAt
		}
		if current.cluster.SampleIntent == "" {
			current.cluster.SampleIntent = strings.TrimSpace(intent)
		}
		groups[k] = current
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	out := make([]CapabilityGapCluster, 0, len(groups))
	for _, v := range groups {
		out = append(out, v.cluster)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Count == out[j].Count {
			return out[i].LastSeenAt.After(out[j].LastSeenAt)
		}
		return out[i].Count > out[j].Count
	})
	if len(out) > limit {
		out = out[:limit]
	}
	return out, nil
}

func normalizeIntentKey(intent string) string {
	intent = strings.ToLower(strings.TrimSpace(intent))
	if intent == "" {
		return "empty"
	}
	intent = strings.ReplaceAll(intent, " ", "")
	intent = strings.ReplaceAll(intent, "\t", "")
	intent = strings.ReplaceAll(intent, "\n", "")
	intent = strings.ReplaceAll(intent, "\r", "")
	b := strings.Builder{}
	for _, r := range intent {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || (r >= 0x4e00 && r <= 0x9fff) {
			b.WriteRune(r)
		}
	}
	normalized := b.String()
	if normalized == "" {
		return "empty"
	}
	runes := []rune(normalized)
	if len(runes) > 80 {
		normalized = string(runes[:80])
	}
	return normalized
}

func (s *SQLiteStore) migrate() error {
	stmt := `
	CREATE TABLE IF NOT EXISTS messages (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		chat_id TEXT NOT NULL,
		user_id TEXT NOT NULL,
		role TEXT NOT NULL,
		content TEXT NOT NULL,
		created_at TIMESTAMP NOT NULL
	);
	CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
	CREATE TABLE IF NOT EXISTS capability_gaps (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		chat_id TEXT NOT NULL,
		user_id TEXT NOT NULL,
		intent TEXT NOT NULL,
		reason TEXT NOT NULL,
		created_at TIMESTAMP NOT NULL
	);
	CREATE INDEX IF NOT EXISTS idx_capability_gaps_created_at ON capability_gaps(created_at);
	`
	if _, err := s.db.Exec(stmt); err != nil {
		return fmt.Errorf("migrate schema: %w", err)
	}
	if s.log != nil {
		s.log.Infof("sqlite schema migration completed")
	}
	return nil
}