MuyueWorkspace/internal/memory/store.go

package memory

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

type MemoryType string

const (
	TypePreference MemoryType = "preference"
	TypeContext    MemoryType = "context"
	TypeSummary    MemoryType = "summary"
	TypeFact       MemoryType = "fact"
	TypePattern    MemoryType = "pattern"
)

type Memory struct {
	ID          string     `json:"id"`
	Type        MemoryType `json:"type"`
	Key         string     `json:"key"`
	Content     string     `json:"content"`
	Tags        string     `json:"tags,omitempty"`
	Source      string     `json:"source,omitempty"`
	Confidence  float64    `json:"confidence,omitempty"`
	AccessCount int        `json:"access_count"`
	CreatedAt   time.Time  `json:"created_at"`
	UpdatedAt   time.Time  `json:"updated_at"`
}

type Store struct {
	db   *sql.DB
	path string
	mu   sync.RWMutex
}

func NewStore() (*Store, error) {
	dbPath, err := dbPath()
	if err != nil {
		return nil, fmt.Errorf("get db path: %w", err)
	}

	if err := os.MkdirAll(filepath.Dir(dbPath), 0755); err != nil {
		return nil, fmt.Errorf("create memory dir: %w", err)
	}

	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, fmt.Errorf("open memory db: %w", err)
	}

	db.SetMaxOpenConns(1)

	s := &Store{db: db, path: dbPath}
	if err := s.migrate(); err != nil {
		db.Close()
		return nil, fmt.Errorf("migrate: %w", err)
	}

	return s, nil
}

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

func (s *Store) Store(m *Memory) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if m.ID == "" {
		m.ID = generateID()
	}

	now := time.Now()
	if m.CreatedAt.IsZero() {
		m.CreatedAt = now
	}
	m.UpdatedAt = now

	_, err := s.db.Exec(`
		INSERT INTO memories (id, type, key, content, tags, source, confidence, access_count, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(id) DO UPDATE SET
			type = excluded.type,
			key = excluded.key,
			content = excluded.content,
			tags = excluded.tags,
			source = excluded.source,
			confidence = excluded.confidence,
			access_count = excluded.access_count,
			updated_at = excluded.updated_at
	`, m.ID, string(m.Type), m.Key, m.Content, m.Tags, m.Source, m.Confidence, m.AccessCount, m.CreatedAt, m.UpdatedAt)

	return err
}

func (s *Store) Get(id string) (*Memory, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	m := &Memory{}
	err := s.db.QueryRow(`
		SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
		FROM memories WHERE id = ?
	`, id).Scan(&m.ID, &m.Type, &m.Key, &m.Content, &m.Tags, &m.Source, &m.Confidence, &m.AccessCount, &m.CreatedAt, &m.UpdatedAt)
	if err == nil {
		s.incrementAccess(id)
	}
	return m, err
}

func (s *Store) Delete(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, err := s.db.Exec(`DELETE FROM memories WHERE id = ?`, id)
	return err
}

func (s *Store) List(memType MemoryType, limit, offset int) ([]Memory, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if limit <= 0 {
		limit = 50
	}
	if limit > 200 {
		limit = 200
	}

	var rows *sql.Rows
	var err error

	if memType != "" {
		rows, err = s.db.Query(`
			SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
			FROM memories WHERE type = ?
			ORDER BY updated_at DESC LIMIT ? OFFSET ?
		`, string(memType), limit, offset)
	} else {
		rows, err = s.db.Query(`
			SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
			FROM memories ORDER BY updated_at DESC LIMIT ? OFFSET ?
		`, limit, offset)
	}

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanMemories(rows)
}

func (s *Store) Count() (int, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var count int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM memories`).Scan(&count)
	return count, err
}

func (s *Store) incrementAccess(id string) {
	go func() {
		s.db.Exec(`UPDATE memories SET access_count = access_count + 1 WHERE id = ?`, id)
	}()
}

func (s *Store) migrate() error {
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS memories (
			id TEXT PRIMARY KEY,
			type TEXT NOT NULL,
			key TEXT NOT NULL,
			content TEXT NOT NULL,
			tags TEXT DEFAULT '',
			source TEXT DEFAULT '',
			confidence REAL DEFAULT 0.5,
			access_count INTEGER DEFAULT 0,
			created_at DATETIME NOT NULL,
			updated_at DATETIME NOT NULL
		)
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(type)
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE INDEX IF NOT EXISTS idx_memories_key ON memories(key)
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5(
			key, content, tags,
			content=memories,
			content_rowid=rowid
		)
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE TRIGGER IF NOT EXISTS memories_ai AFTER INSERT ON memories BEGIN
			INSERT INTO memories_fts(rowid, key, content, tags)
			VALUES (new.rowid, new.key, new.content, new.tags);
		END
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE TRIGGER IF NOT EXISTS memories_ad AFTER DELETE ON memories BEGIN
			INSERT INTO memories_fts(memories_fts, rowid, key, content, tags)
			VALUES ('delete', old.rowid, old.key, old.content, old.tags);
		END
	`)
	if err != nil {
		return err
	}

	_, err = s.db.Exec(`
		CREATE TRIGGER IF NOT EXISTS memories_au AFTER UPDATE ON memories BEGIN
			INSERT INTO memories_fts(memories_fts, rowid, key, content, tags)
			VALUES ('delete', old.rowid, old.key, old.content, old.tags);
			INSERT INTO memories_fts(rowid, key, content, tags)
			VALUES (new.rowid, new.key, new.content, new.tags);
		END
	`)

	return err
}

func scanMemories(rows *sql.Rows) ([]Memory, error) {
	var memories []Memory
	for rows.Next() {
		var m Memory
		err := rows.Scan(&m.ID, &m.Type, &m.Key, &m.Content, &m.Tags, &m.Source, &m.Confidence, &m.AccessCount, &m.CreatedAt, &m.UpdatedAt)
		if err != nil {
			return memories, err
		}
		memories = append(memories, m)
	}
	return memories, nil
}

func dbPath() (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	return filepath.Join(home, ".muyue", "memory", "memories.db"), nil
}

func generateID() string {
	return fmt.Sprintf("mem_%d", time.Now().UnixNano())
}