AIRouter/internal/session/session.go

package session

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/edgeai/gateway/pkg/api"
	_ "github.com/mattn/go-sqlite3"
)

// Session represents a conversation session.
type Session struct {
	ID            string
	ApplicationID string
	TenantID      string
	UserID        string
	Messages      []api.Message
	Config        map[string]any
	CreatedAt     time.Time
	LastActive    time.Time
}

// Store manages session persistence with SQLite.
type Store struct {
	mu sync.RWMutex
	db *sql.DB
}

// NewStore creates a new session store.
func NewStore(dbPath string) (*Store, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("open session db: %w", err)
	}

	if err := initSessionDB(db); err != nil {
		return nil, fmt.Errorf("init session db: %w", err)
	}

	return &Store{db: db}, nil
}

func initSessionDB(db *sql.DB) error {
	schema := `
	CREATE TABLE IF NOT EXISTS sessions (
		id TEXT PRIMARY KEY,
		application_id TEXT NOT NULL,
		tenant_id TEXT NOT NULL,
		user_id TEXT,
		messages TEXT NOT NULL DEFAULT '[]',
		config TEXT NOT NULL DEFAULT '{}',
		created_at TEXT NOT NULL,
		last_active TEXT NOT NULL
	);
	CREATE INDEX IF NOT EXISTS idx_sessions_app ON sessions(application_id);
	CREATE INDEX IF NOT EXISTS idx_sessions_tenant ON sessions(tenant_id);
	CREATE INDEX IF NOT EXISTS idx_sessions_last_active ON sessions(last_active);`
	_, err := db.Exec(schema)
	return err
}

// Create creates a new session.
func (s *Store) Create(id, appID, tenantID, userID string, config map[string]any) (*Session, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	session := &Session{
		ID:            id,
		ApplicationID: appID,
		TenantID:      tenantID,
		UserID:        userID,
		Messages:      []api.Message{},
		Config:        config,
		CreatedAt:     now,
		LastActive:    now,
	}

	configJSON, _ := json.Marshal(config)
	msgsJSON, _ := json.Marshal(session.Messages)

	_, err := s.db.Exec(
		`INSERT INTO sessions (id, application_id, tenant_id, user_id, messages, config, created_at, last_active)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
		id, appID, tenantID, userID, string(msgsJSON), string(configJSON), now.Format(time.RFC3339), now.Format(time.RFC3339),
	)
	if err != nil {
		return nil, fmt.Errorf("insert session: %w", err)
	}

	return session, nil
}

// Get retrieves a session by ID.
func (s *Store) Get(id string) (*Session, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var (
		appID, tenantID, userID, msgsJSON, configJSON, createdAt, lastActive string
	)

	err := s.db.QueryRow(
		`SELECT application_id, tenant_id, user_id, messages, config, created_at, last_active FROM sessions WHERE id = ?`,
		id,
	).Scan(&appID, &tenantID, &userID, &msgsJSON, &configJSON, &createdAt, &lastActive)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("query session: %w", err)
	}

	session := &Session{
		ID:            id,
		ApplicationID: appID,
		TenantID:      tenantID,
		UserID:        userID,
		CreatedAt:     parseTime(createdAt),
		LastActive:    parseTime(lastActive),
	}
	json.Unmarshal([]byte(msgsJSON), &session.Messages)
	json.Unmarshal([]byte(configJSON), &session.Config)

	return session, nil
}

// GetForApp 检索会话并校验 application_id 是否匹配。
// 如果会话不属于该 app，返回 nil（权限隔离）。
func (s *Store) GetForApp(id, appID string) (*Session, error) {
	sess, err := s.Get(id)
	if err != nil {
		return nil, err
	}
	if sess == nil {
		return nil, nil
	}
	if sess.ApplicationID != appID {
		return nil, nil
	}
	return sess, nil
}

// DeleteForApp 删除会话并校验 application_id 是否匹配。
// 如果会话不属于该 app，返回错误。
func (s *Store) DeleteForApp(id, appID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	var existingAppID string
	err := s.db.QueryRow(`SELECT application_id FROM sessions WHERE id = ?`, id).Scan(&existingAppID)
	if err == sql.ErrNoRows {
		return fmt.Errorf("session not found: %s", id)
	}
	if err != nil {
		return fmt.Errorf("query session: %w", err)
	}
	if existingAppID != appID {
		return fmt.Errorf("session does not belong to application: %s", appID)
	}

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

// AddMessage 追加消息到会话并更新 last_active。
// 当消息数超过 maxMessages 时，自动裁剪最旧的消息。
// 注意：不能调用 s.Get()，因为 Get 会获取读锁，而此处已持有写锁，会导致死锁。
func (s *Store) AddMessage(id string, msg api.Message) error {
	return s.AddMessageWithLimit(id, msg, 0)
}

// AddMessageWithLimit 追加消息到会话，当 maxMessages > 0 时限制消息总数。
// 超出上限时自动裁剪最旧的消息，保留最新的 maxMessages 条。
func (s *Store) AddMessageWithLimit(id string, msg api.Message, maxMessages int) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	var msgsJSON string
	err := s.db.QueryRow(`SELECT messages FROM sessions WHERE id = ?`, id).Scan(&msgsJSON)
	if err == sql.ErrNoRows {
		return fmt.Errorf("session not found: %s", id)
	}
	if err != nil {
		return fmt.Errorf("query session messages: %w", err)
	}

	var msgs []api.Message
	if err := json.Unmarshal([]byte(msgsJSON), &msgs); err != nil {
		return fmt.Errorf("unmarshal session messages: %w", err)
	}

	msgs = append(msgs, msg)

	// 消息数上限保护：裁剪最旧消息
	if maxMessages > 0 && len(msgs) > maxMessages {
		msgs = msgs[len(msgs)-maxMessages:]
	}

	newMsgsJSON, _ := json.Marshal(msgs)
	now := time.Now().Format(time.RFC3339)

	_, err = s.db.Exec(
		`UPDATE sessions SET messages = ?, last_active = ? WHERE id = ?`,
		string(newMsgsJSON), now, id,
	)
	return err
}

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

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

// ListByApp 列出指定 app 的会话（分页）。
func (s *Store) ListByApp(appID string, limit, offset int) ([]*Session, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if limit <= 0 || limit > 100 {
		limit = 20
	}

	rows, err := s.db.Query(
		`SELECT id, application_id, tenant_id, user_id, messages, config, created_at, last_active
		 FROM sessions WHERE application_id = ? ORDER BY last_active DESC LIMIT ? OFFSET ?`,
		appID, limit, offset,
	)
	if err != nil {
		return nil, fmt.Errorf("list sessions: %w", err)
	}
	defer rows.Close()

	var sessions []*Session
	for rows.Next() {
		var (
			sid, sAppID, sTenantID, sUserID, msgsJSON, configJSON, createdAt, lastActive string
		)
		if err := rows.Scan(&sid, &sAppID, &sTenantID, &sUserID, &msgsJSON, &configJSON, &createdAt, &lastActive); err != nil {
			return nil, err
		}
		sess := &Session{
			ID:            sid,
			ApplicationID: sAppID,
			TenantID:      sTenantID,
			UserID:        sUserID,
			CreatedAt:     parseTime(createdAt),
			LastActive:    parseTime(lastActive),
		}
		json.Unmarshal([]byte(msgsJSON), &sess.Messages)
		json.Unmarshal([]byte(configJSON), &sess.Config)
		sessions = append(sessions, sess)
	}
	return sessions, rows.Err()
}

// Close closes the database connection.
func (s *Store) Close() error {
	return s.db.Close()
}

// CleanupExpired 清理超过 TTL 未活跃的会话。
// ttlMinutes 为会话空闲超时（分钟），超过此时间未活跃的会话将被删除。
// 返回删除的会话数量。
func (s *Store) CleanupExpired(ttlMinutes int) (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	cutoff := time.Now().Add(-time.Duration(ttlMinutes) * time.Minute).Format(time.RFC3339)
	result, err := s.db.Exec(`DELETE FROM sessions WHERE last_active < ?`, cutoff)
	if err != nil {
		return 0, fmt.Errorf("cleanup expired sessions: %w", err)
	}
	n, _ := result.RowsAffected()
	return n, nil
}

// StartCleanupTask 启动后台定时清理任务。
// ttlMinutes 为会话空闲超时（分钟），intervalMinutes 为清理间隔（分钟）。
// onCleanup 为清理完成后的回调（可传 nil），参数为清理的会话数量。
// 返回停止函数，调用以停止清理任务。
func (s *Store) StartCleanupTask(ttlMinutes, intervalMinutes int, onCleanup func(deleted int64)) func() {
	ticker := time.NewTicker(time.Duration(intervalMinutes) * time.Minute)
	stop := make(chan struct{})

	go func() {
		for {
			select {
			case <-ticker.C:
				if n, err := s.CleanupExpired(ttlMinutes); err == nil && n > 0 {
					if onCleanup != nil {
						onCleanup(n)
					}
				}
			case <-stop:
				ticker.Stop()
				return
			}
		}
	}()

	return func() { close(stop) }
}

func parseTime(s string) time.Time {
	t, _ := time.Parse(time.RFC3339, s)
	return t
}