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
}

// AddMessage appends a message to the session and updates last_active.
func (s *Store) AddMessage(id string, msg api.Message) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	session, err := s.Get(id)
	if err != nil {
		return err
	}
	if session == nil {
		return fmt.Errorf("session not found: %s", id)
	}

	session.Messages = append(session.Messages, msg)
	msgsJSON, _ := json.Marshal(session.Messages)
	now := time.Now().Format(time.RFC3339)

	_, err = s.db.Exec(
		`UPDATE sessions SET messages = ?, last_active = ? WHERE id = ?`,
		string(msgsJSON), 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
}

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

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