codex-telegram-bot/internal/store/store.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	_ "modernc.org/sqlite"
)

type Store struct {
	db *sql.DB
}

type AllowedUser struct {
	TelegramUserID int64
	Username       string
	Notes          string
	AddedAt        string
}

type Workspace struct {
	ID        int64
	Path      string
	Label     string
	IsDefault bool
	CreatedAt string
}

type Session struct {
	TelegramUserID    int64
	ActiveThreadID    int64
	ActiveWorkspaceID int64
	Model             string
	ReasoningEffort   string
	Sandbox           string
	ActiveTurnID      string
	SettingsChatID    int64
	SettingsMessageID int
	UpdatedAt         string
}

type Thread struct {
	ID             int64
	TelegramUserID int64
	CodexThreadID  string
	WorkspaceID    int64
	Title          string
	Archived       bool
	CreatedAt      string
	UpdatedAt      string
}

type PendingApproval struct {
	ID             int64
	TelegramUserID int64
	CodexRequestID string
	CodexThreadID  string
	TurnID         string
	ItemID         string
	Kind           string
	PayloadJSON    string
	MessageChatID  int64
	MessageID      int
	Status         string
	CreatedAt      string
	ResolvedAt     string
}

func Open(ctx context.Context, path string) (*Store, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(1)

	s := &Store{db: db}
	if err := s.configure(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	if err := s.migrate(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	return s, nil
}

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

func (s *Store) configure(ctx context.Context) error {
	statements := []string{
		"PRAGMA foreign_keys = ON",
		"PRAGMA journal_mode = WAL",
		"PRAGMA busy_timeout = 5000",
	}
	for _, statement := range statements {
		if _, err := s.db.ExecContext(ctx, statement); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) migrate(ctx context.Context) error {
	if _, err := s.db.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS schema_migrations (
	version INTEGER PRIMARY KEY,
	name TEXT NOT NULL,
	applied_at TEXT NOT NULL DEFAULT (datetime('now'))
)`); err != nil {
		return err
	}

	var current int
	if err := s.db.QueryRowContext(ctx, "SELECT COALESCE(MAX(version), 0) FROM schema_migrations").Scan(&current); err != nil {
		return err
	}
	for i := current; i < len(migrations); i++ {
		version := i + 1
		tx, err := s.db.BeginTx(ctx, nil)
		if err != nil {
			return err
		}
		if _, err := tx.ExecContext(ctx, migrations[i].SQL); err != nil {
			_ = tx.Rollback()
			return fmt.Errorf("migration %d %s: %w", version, migrations[i].Name, err)
		}
		if _, err := tx.ExecContext(ctx, "INSERT INTO schema_migrations (version, name) VALUES (?, ?)", version, migrations[i].Name); err != nil {
			_ = tx.Rollback()
			return err
		}
		if err := tx.Commit(); err != nil {
			return err
		}
	}
	return nil
}

func ValidateWorkspacePath(path string) (string, error) {
	if path == "" {
		return "", errors.New("workspace path is required")
	}
	if strings.ContainsRune(path, 0) {
		return "", errors.New("workspace path contains a NUL byte")
	}
	if !filepath.IsAbs(path) {
		return "", errors.New("workspace path must be absolute")
	}
	clean := filepath.Clean(path)
	if clean == string(filepath.Separator) {
		return "", errors.New("workspace path cannot be filesystem root")
	}
	return clean, nil
}

func (s *Store) IsAllowed(ctx context.Context, telegramUserID int64) (bool, error) {
	var exists int
	err := s.db.QueryRowContext(ctx, "SELECT 1 FROM allowed_users WHERE telegram_user_id = ?", telegramUserID).Scan(&exists)
	if errors.Is(err, sql.ErrNoRows) {
		return false, nil
	}
	return err == nil, err
}

func (s *Store) AddAllowedUser(ctx context.Context, telegramUserID int64, username, notes string) error {
	_, err := s.db.ExecContext(ctx, `
INSERT INTO allowed_users (telegram_user_id, username, notes)
VALUES (?, ?, ?)
ON CONFLICT(telegram_user_id) DO UPDATE SET username = excluded.username, notes = excluded.notes`,
		telegramUserID, username, notes)
	return err
}

func (s *Store) RemoveAllowedUser(ctx context.Context, telegramUserID int64) error {
	_, err := s.db.ExecContext(ctx, "DELETE FROM allowed_users WHERE telegram_user_id = ?", telegramUserID)
	return err
}

func (s *Store) ListAllowedUsers(ctx context.Context) ([]AllowedUser, error) {
	rows, err := s.db.QueryContext(ctx, "SELECT telegram_user_id, username, notes, added_at FROM allowed_users ORDER BY telegram_user_id")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var users []AllowedUser
	for rows.Next() {
		var user AllowedUser
		if err := rows.Scan(&user.TelegramUserID, &user.Username, &user.Notes, &user.AddedAt); err != nil {
			return nil, err
		}
		users = append(users, user)
	}
	return users, rows.Err()
}

func (s *Store) AddWorkspace(ctx context.Context, path, label string, isDefault bool) (Workspace, error) {
	clean, err := ValidateWorkspacePath(path)
	if err != nil {
		return Workspace{}, err
	}
	if label == "" {
		label = filepath.Base(clean)
	}

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return Workspace{}, err
	}
	if isDefault {
		if _, err := tx.ExecContext(ctx, "UPDATE workspaces SET is_default = 0"); err != nil {
			_ = tx.Rollback()
			return Workspace{}, err
		}
	}
	result, err := tx.ExecContext(ctx, `
INSERT INTO workspaces (path, label, is_default)
VALUES (?, ?, ?)
ON CONFLICT(path) DO UPDATE SET label = excluded.label, is_default = excluded.is_default`,
		clean, label, boolInt(isDefault))
	if err != nil {
		_ = tx.Rollback()
		return Workspace{}, err
	}
	if err := tx.Commit(); err != nil {
		return Workspace{}, err
	}

	id, _ := result.LastInsertId()
	if id == 0 {
		return s.GetWorkspaceByPath(ctx, clean)
	}
	return s.GetWorkspaceByPath(ctx, clean)
}

func (s *Store) ListWorkspaces(ctx context.Context) ([]Workspace, error) {
	rows, err := s.db.QueryContext(ctx, "SELECT id, path, label, is_default, created_at FROM workspaces ORDER BY is_default DESC, label, path")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var workspaces []Workspace
	for rows.Next() {
		workspace, err := scanWorkspace(rows)
		if err != nil {
			return nil, err
		}
		workspaces = append(workspaces, workspace)
	}
	return workspaces, rows.Err()
}

func (s *Store) GetWorkspaceByID(ctx context.Context, id int64) (Workspace, error) {
	row := s.db.QueryRowContext(ctx, "SELECT id, path, label, is_default, created_at FROM workspaces WHERE id = ?", id)
	return scanWorkspace(row)
}

func (s *Store) GetWorkspaceByPath(ctx context.Context, path string) (Workspace, error) {
	row := s.db.QueryRowContext(ctx, "SELECT id, path, label, is_default, created_at FROM workspaces WHERE path = ?", path)
	return scanWorkspace(row)
}

func (s *Store) DefaultWorkspace(ctx context.Context) (Workspace, error) {
	row := s.db.QueryRowContext(ctx, "SELECT id, path, label, is_default, created_at FROM workspaces ORDER BY is_default DESC, id ASC LIMIT 1")
	return scanWorkspace(row)
}

func (s *Store) GetOrCreateSession(ctx context.Context, telegramUserID int64, defaultModel, defaultSandbox string) (Session, error) {
	if defaultSandbox == "" {
		defaultSandbox = "workspace-write"
	}
	_, err := s.db.ExecContext(ctx, `
INSERT INTO sessions (telegram_user_id, model, sandbox)
VALUES (?, ?, ?)
ON CONFLICT(telegram_user_id) DO NOTHING`, telegramUserID, defaultModel, defaultSandbox)
	if err != nil {
		return Session{}, err
	}
	return s.GetSession(ctx, telegramUserID)
}

func (s *Store) GetSession(ctx context.Context, telegramUserID int64) (Session, error) {
	row := s.db.QueryRowContext(ctx, `
SELECT telegram_user_id, COALESCE(active_thread_id, 0), COALESCE(active_workspace_id, 0), model, COALESCE(reasoning_effort, ''), sandbox, active_turn_id, COALESCE(settings_chat_id, 0), COALESCE(settings_message_id, 0), updated_at
FROM sessions WHERE telegram_user_id = ?`, telegramUserID)
	var session Session
	err := row.Scan(&session.TelegramUserID, &session.ActiveThreadID, &session.ActiveWorkspaceID, &session.Model, &session.ReasoningEffort, &session.Sandbox, &session.ActiveTurnID, &session.SettingsChatID, &session.SettingsMessageID, &session.UpdatedAt)
	return session, err
}

func (s *Store) SetSessionWorkspace(ctx context.Context, telegramUserID, workspaceID int64) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE sessions SET active_workspace_id = ?, updated_at = datetime('now')
WHERE telegram_user_id = ?`, workspaceID, telegramUserID)
	return err
}

func (s *Store) SetSessionModel(ctx context.Context, telegramUserID int64, model string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET model = ?, reasoning_effort = '', updated_at = datetime('now') WHERE telegram_user_id = ?", model, telegramUserID)
	return err
}

func (s *Store) SetSessionReasoningEffort(ctx context.Context, telegramUserID int64, effort string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET reasoning_effort = ?, updated_at = datetime('now') WHERE telegram_user_id = ?", effort, telegramUserID)
	return err
}

func (s *Store) SetSessionSettingsMessage(ctx context.Context, telegramUserID int64, chatID int64, messageID int) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET settings_chat_id = ?, settings_message_id = ?, updated_at = datetime('now') WHERE telegram_user_id = ?", chatID, messageID, telegramUserID)
	return err
}

func (s *Store) SetSessionSandbox(ctx context.Context, telegramUserID int64, sandbox string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET sandbox = ?, updated_at = datetime('now') WHERE telegram_user_id = ?", sandbox, telegramUserID)
	return err
}

func (s *Store) SetActiveThread(ctx context.Context, telegramUserID, threadID int64) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE sessions SET active_thread_id = ?, active_turn_id = '', updated_at = datetime('now')
WHERE telegram_user_id = ?`, threadID, telegramUserID)
	return err
}

func (s *Store) SetActiveTurn(ctx context.Context, telegramUserID int64, turnID string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET active_turn_id = ?, updated_at = datetime('now') WHERE telegram_user_id = ?", turnID, telegramUserID)
	return err
}

func (s *Store) ClearActiveTurns(ctx context.Context) error {
	_, err := s.db.ExecContext(ctx, "UPDATE sessions SET active_turn_id = '', updated_at = datetime('now') WHERE active_turn_id <> ''")
	return err
}

func (s *Store) CreateThread(ctx context.Context, telegramUserID int64, codexThreadID string, workspaceID int64, title string) (Thread, error) {
	result, err := s.db.ExecContext(ctx, `
INSERT INTO threads (telegram_user_id, codex_thread_id, workspace_id, title)
VALUES (?, ?, ?, ?)`, telegramUserID, codexThreadID, workspaceID, title)
	if err != nil {
		return Thread{}, err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return Thread{}, err
	}
	return s.GetThreadByID(ctx, telegramUserID, id)
}

func (s *Store) GetThreadByID(ctx context.Context, telegramUserID, id int64) (Thread, error) {
	row := s.db.QueryRowContext(ctx, `
SELECT id, telegram_user_id, codex_thread_id, workspace_id, title, archived, created_at, updated_at
FROM threads WHERE telegram_user_id = ? AND id = ?`, telegramUserID, id)
	return scanThread(row)
}

func (s *Store) GetThreadByCodexID(ctx context.Context, codexThreadID string) (Thread, error) {
	row := s.db.QueryRowContext(ctx, `
SELECT id, telegram_user_id, codex_thread_id, workspace_id, title, archived, created_at, updated_at
FROM threads WHERE codex_thread_id = ?`, codexThreadID)
	return scanThread(row)
}

func (s *Store) ListThreads(ctx context.Context, telegramUserID int64, includeArchived bool) ([]Thread, error) {
	return s.ListThreadsPage(ctx, telegramUserID, includeArchived, 20, 0)
}

func (s *Store) ListThreadsPage(ctx context.Context, telegramUserID int64, includeArchived bool, limit, offset int) ([]Thread, error) {
	if limit <= 0 {
		limit = 20
	}
	if offset < 0 {
		offset = 0
	}
	query := `
SELECT id, telegram_user_id, codex_thread_id, workspace_id, title, archived, created_at, updated_at
FROM threads WHERE telegram_user_id = ?`
	args := []any{telegramUserID}
	if !includeArchived {
		query += " AND archived = 0"
	}
	query += " ORDER BY updated_at DESC, id DESC LIMIT ? OFFSET ?"
	args = append(args, limit, offset)

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var threads []Thread
	for rows.Next() {
		thread, err := scanThread(rows)
		if err != nil {
			return nil, err
		}
		threads = append(threads, thread)
	}
	return threads, rows.Err()
}

func (s *Store) ArchiveThread(ctx context.Context, telegramUserID, id int64) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE threads SET archived = 1, updated_at = datetime('now')
WHERE telegram_user_id = ? AND id = ?`, telegramUserID, id)
	return err
}

func (s *Store) TouchThread(ctx context.Context, codexThreadID string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE threads SET updated_at = datetime('now') WHERE codex_thread_id = ?", codexThreadID)
	return err
}

func (s *Store) RenameThread(ctx context.Context, telegramUserID, id int64, title string) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE threads SET title = ?, updated_at = datetime('now')
WHERE telegram_user_id = ? AND id = ?`, title, telegramUserID, id)
	return err
}

func (s *Store) RenameThreadByCodexID(ctx context.Context, codexThreadID, title string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE threads SET title = ?, updated_at = datetime('now') WHERE codex_thread_id = ?", title, codexThreadID)
	return err
}

func (s *Store) SyncThreadTitle(ctx context.Context, telegramUserID, id int64, title string) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE threads SET title = ?
WHERE telegram_user_id = ? AND id = ?`, title, telegramUserID, id)
	return err
}

func (s *Store) SyncThreadTitleByCodexID(ctx context.Context, codexThreadID, title string) error {
	_, err := s.db.ExecContext(ctx, "UPDATE threads SET title = ? WHERE codex_thread_id = ?", title, codexThreadID)
	return err
}

func (s *Store) SyncThreadWorkspace(ctx context.Context, telegramUserID, id, workspaceID int64) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE threads SET workspace_id = ?
WHERE telegram_user_id = ? AND id = ?`, workspaceID, telegramUserID, id)
	return err
}

func (s *Store) UpsertPendingApproval(ctx context.Context, approval PendingApproval) (PendingApproval, error) {
	_, err := s.db.ExecContext(ctx, `
INSERT INTO pending_approvals (
	telegram_user_id, codex_request_id, codex_thread_id, turn_id, item_id, kind, payload_json,
	message_chat_id, message_id, status
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')
ON CONFLICT(telegram_user_id, codex_request_id) DO UPDATE SET
	payload_json = excluded.payload_json,
	message_chat_id = CASE WHEN pending_approvals.status = 'pending' THEN excluded.message_chat_id ELSE pending_approvals.message_chat_id END,
	message_id = CASE WHEN pending_approvals.status = 'pending' THEN excluded.message_id ELSE pending_approvals.message_id END`,
		approval.TelegramUserID, approval.CodexRequestID, approval.CodexThreadID, approval.TurnID,
		approval.ItemID, approval.Kind, approval.PayloadJSON, approval.MessageChatID, approval.MessageID)
	if err != nil {
		return PendingApproval{}, err
	}
	return s.GetPendingApprovalByRequest(ctx, approval.TelegramUserID, approval.CodexRequestID)
}

func (s *Store) UpdatePendingApprovalMessage(ctx context.Context, id int64, chatID int64, messageID int) error {
	_, err := s.db.ExecContext(ctx, "UPDATE pending_approvals SET message_chat_id = ?, message_id = ? WHERE id = ?", chatID, messageID, id)
	return err
}

func (s *Store) GetPendingApproval(ctx context.Context, telegramUserID, id int64) (PendingApproval, error) {
	row := s.db.QueryRowContext(ctx, `
SELECT id, telegram_user_id, codex_request_id, codex_thread_id, turn_id, item_id, kind, payload_json,
       message_chat_id, message_id, status, created_at, COALESCE(resolved_at, '')
FROM pending_approvals WHERE telegram_user_id = ? AND id = ?`, telegramUserID, id)
	return scanPendingApproval(row)
}

func (s *Store) GetPendingApprovalByRequest(ctx context.Context, telegramUserID int64, requestID string) (PendingApproval, error) {
	row := s.db.QueryRowContext(ctx, `
SELECT id, telegram_user_id, codex_request_id, codex_thread_id, turn_id, item_id, kind, payload_json,
       message_chat_id, message_id, status, created_at, COALESCE(resolved_at, '')
FROM pending_approvals WHERE telegram_user_id = ? AND codex_request_id = ?`, telegramUserID, requestID)
	return scanPendingApproval(row)
}

func (s *Store) ResolvePendingApproval(ctx context.Context, telegramUserID, id int64, status string) error {
	_, err := s.db.ExecContext(ctx, `
UPDATE pending_approvals SET status = ?, resolved_at = datetime('now')
WHERE telegram_user_id = ? AND id = ? AND status = 'pending'`, status, telegramUserID, id)
	return err
}

func (s *Store) Audit(ctx context.Context, telegramUserID int64, action, details string) error {
	_, err := s.db.ExecContext(ctx, "INSERT INTO audit_log (telegram_user_id, action, details) VALUES (?, ?, ?)", telegramUserID, action, details)
	return err
}

type scanner interface {
	Scan(dest ...any) error
}

func scanWorkspace(row scanner) (Workspace, error) {
	var workspace Workspace
	var isDefault int
	if err := row.Scan(&workspace.ID, &workspace.Path, &workspace.Label, &isDefault, &workspace.CreatedAt); err != nil {
		return Workspace{}, err
	}
	workspace.IsDefault = isDefault != 0
	return workspace, nil
}

func scanThread(row scanner) (Thread, error) {
	var thread Thread
	var archived int
	if err := row.Scan(&thread.ID, &thread.TelegramUserID, &thread.CodexThreadID, &thread.WorkspaceID, &thread.Title, &archived, &thread.CreatedAt, &thread.UpdatedAt); err != nil {
		return Thread{}, err
	}
	thread.Archived = archived != 0
	return thread, nil
}

func scanPendingApproval(row scanner) (PendingApproval, error) {
	var approval PendingApproval
	if err := row.Scan(&approval.ID, &approval.TelegramUserID, &approval.CodexRequestID, &approval.CodexThreadID,
		&approval.TurnID, &approval.ItemID, &approval.Kind, &approval.PayloadJSON, &approval.MessageChatID,
		&approval.MessageID, &approval.Status, &approval.CreatedAt, &approval.ResolvedAt); err != nil {
		return PendingApproval{}, err
	}
	return approval, nil
}

func boolInt(value bool) int {
	if value {
		return 1
	}
	return 0
}