notifier/internal/auth/keystore_db.go

package auth

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"github.com/lib/pq"
	_ "github.com/lib/pq" // PostgreSQL driver
)

// KeyStoreDB provides persistent storage for API keys using PostgreSQL
// It acts as the backend for the in-memory cache
type KeyStoreDB struct {
	db *sql.DB
}

// NewKeyStoreDB creates a new database-backed key store
func NewKeyStoreDB(dbURL string) (*KeyStoreDB, error) {
	db, err := sql.Open("postgres", dbURL)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}

	// Test connection
	if err := db.Ping(); err != nil {
		return nil, fmt.Errorf("failed to ping database: %w", err)
	}

	ks := &KeyStoreDB{db: db}

	// Initialize schema
	if err := ks.initializeSchema(); err != nil {
		return nil, fmt.Errorf("failed to initialize schema: %w", err)
	}

	return ks, nil
}

// initializeSchema creates the necessary tables and indexes if they don't exist
func (ks *KeyStoreDB) initializeSchema() error {
	// Create tables
	tableSchema := `
	-- API Keys table
	CREATE TABLE IF NOT EXISTS api_keys (
		id SERIAL PRIMARY KEY,
		key VARCHAR(255) UNIQUE NOT NULL,
		name VARCHAR(255) NOT NULL,
		client_id VARCHAR(255) NOT NULL,
		roles TEXT[] DEFAULT '{}',
		created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
		last_used_at TIMESTAMP,
		expires_at TIMESTAMP,
		is_active BOOLEAN NOT NULL DEFAULT true,
		rate_limit INTEGER NOT NULL DEFAULT 0,
		created_by VARCHAR(255),
		metadata JSONB DEFAULT '{}'::jsonb
	);

	-- Audit log for key operations
	CREATE TABLE IF NOT EXISTS api_key_audit_log (
		id SERIAL PRIMARY KEY,
		key_id INTEGER NOT NULL REFERENCES api_keys(id),
		action VARCHAR(50) NOT NULL,
		performed_by VARCHAR(255) NOT NULL,
		performed_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
		details JSONB DEFAULT '{}'::jsonb
	);
	`

	if _, err := ks.db.Exec(tableSchema); err != nil {
		return fmt.Errorf("failed to create tables: %w", err)
	}

	// Create indexes separately (PostgreSQL syntax)
	indexSchema := `
	CREATE INDEX IF NOT EXISTS idx_api_keys_key ON api_keys(key);
	CREATE INDEX IF NOT EXISTS idx_api_keys_client_id ON api_keys(client_id);
	CREATE INDEX IF NOT EXISTS idx_api_keys_active ON api_keys(is_active);
	CREATE INDEX IF NOT EXISTS idx_api_keys_expires ON api_keys(expires_at);
	CREATE INDEX IF NOT EXISTS idx_api_key_audit_log_key_id ON api_key_audit_log(key_id);
	CREATE INDEX IF NOT EXISTS idx_api_key_audit_log_performed_at ON api_key_audit_log(performed_at);
	`

	if _, err := ks.db.Exec(indexSchema); err != nil {
		return fmt.Errorf("failed to create indexes: %w", err)
	}

	return nil
}

// SaveKey persists an API key to the database
func (ks *KeyStoreDB) SaveKey(ctx context.Context, key *APIKey, createdBy string) error {
	query := `
	INSERT INTO api_keys (key, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit, created_by)
	VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
	ON CONFLICT (key) DO UPDATE SET
		name = EXCLUDED.name,
		roles = EXCLUDED.roles,
		is_active = EXCLUDED.is_active,
		rate_limit = EXCLUDED.rate_limit,
		last_used_at = COALESCE(EXCLUDED.last_used_at, api_keys.last_used_at)
	`

	_, err := ks.db.ExecContext(ctx, query,
		key.Key,
		key.Name,
		key.ClientID,
		pq.Array(key.Roles), // Convert Go slice to PostgreSQL array
		key.CreatedAt,
		key.LastUsedAt,
		key.ExpiresAt,
		key.IsActive,
		key.RateLimit,
		createdBy,
	)

	if err != nil {
		return fmt.Errorf("failed to save key: %w", err)
	}

	// Log to audit trail
	ks.logAudit(ctx, key.Key, "created", createdBy, map[string]interface{}{
		"client_id": key.ClientID,
		"roles":     key.Roles,
	})

	return nil
}

// GetKey retrieves an API key from the database
func (ks *KeyStoreDB) GetKey(ctx context.Context, keyStr string) (*APIKey, error) {
	query := `
	SELECT key, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit
	FROM api_keys
	WHERE key = $1
	`

	var key APIKey
	var roles []string

	err := ks.db.QueryRowContext(ctx, query, keyStr).Scan(
		&key.Key,
		&key.Name,
		&key.ClientID,
		&roles,
		&key.CreatedAt,
		&key.LastUsedAt,
		&key.ExpiresAt,
		&key.IsActive,
		&key.RateLimit,
	)

	if err == sql.ErrNoRows {
		return nil, ErrKeyNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get key: %w", err)
	}

	key.Roles = roles
	return &key, nil
}

// ListKeys retrieves all API keys for a client
func (ks *KeyStoreDB) ListKeys(ctx context.Context, clientID string) ([]*APIKey, error) {
	query := `
	SELECT key, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit
	FROM api_keys
	WHERE client_id = $1 AND is_active = true
	ORDER BY created_at DESC
	`

	rows, err := ks.db.QueryContext(ctx, query, clientID)
	if err != nil {
		return nil, fmt.Errorf("failed to list keys: %w", err)
	}
	defer rows.Close()

	var keys []*APIKey
	for rows.Next() {
		var key APIKey
		var roles []string

		err := rows.Scan(
			&key.Key,
			&key.Name,
			&key.ClientID,
			&roles,
			&key.CreatedAt,
			&key.LastUsedAt,
			&key.ExpiresAt,
			&key.IsActive,
			&key.RateLimit,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan key: %w", err)
		}

		key.Roles = roles
		keys = append(keys, &key)
	}

	return keys, rows.Err()
}

// DeactivateKey disables an API key
func (ks *KeyStoreDB) DeactivateKey(ctx context.Context, keyStr string, deactivatedBy string) error {
	query := `UPDATE api_keys SET is_active = false WHERE key = $1`

	result, err := ks.db.ExecContext(ctx, query, keyStr)
	if err != nil {
		return fmt.Errorf("failed to deactivate key: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return err
	}

	if rows == 0 {
		return ErrKeyNotFound
	}

	ks.logAudit(ctx, keyStr, "deactivated", deactivatedBy, nil)
	return nil
}

// UpdateLastUsed updates the last_used_at timestamp
func (ks *KeyStoreDB) UpdateLastUsed(ctx context.Context, keyStr string) error {
	query := `UPDATE api_keys SET last_used_at = CURRENT_TIMESTAMP WHERE key = $1`

	_, err := ks.db.ExecContext(ctx, query, keyStr)
	if err != nil {
		return fmt.Errorf("failed to update last used: %w", err)
	}

	return nil
}

// LoadAllKeys loads all active keys into memory for caching
func (ks *KeyStoreDB) LoadAllKeys(ctx context.Context) ([]*APIKey, error) {
	query := `
	SELECT key, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit
	FROM api_keys
	WHERE is_active = true AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
	`

	rows, err := ks.db.QueryContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to load keys: %w", err)
	}
	defer rows.Close()

	var keys []*APIKey
	for rows.Next() {
		var key APIKey
		var roles []string

		err := rows.Scan(
			&key.Key,
			&key.Name,
			&key.ClientID,
			&roles,
			&key.CreatedAt,
			&key.LastUsedAt,
			&key.ExpiresAt,
			&key.IsActive,
			&key.RateLimit,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan key: %w", err)
		}

		key.Roles = roles
		keys = append(keys, &key)
	}

	return keys, rows.Err()
}

// logAudit logs a key operation to the audit trail
func (ks *KeyStoreDB) logAudit(ctx context.Context, keyStr string, action string, performedBy string, details map[string]interface{}) {
	// Get key ID
	var keyID int
	err := ks.db.QueryRowContext(ctx, "SELECT id FROM api_keys WHERE key = $1", keyStr).Scan(&keyID)
	if err != nil {
		return // Silently fail audit logging
	}

	// Log the action
	detailsJSON := "{}"
	if len(details) > 0 {
		// Simple JSON encoding (could use jsonb package for robustness)
		detailsJSON = fmt.Sprintf(`{"event": "%s"}`, action)
	}

	query := `
	INSERT INTO api_key_audit_log (key_id, action, performed_by, details)
	VALUES ($1, $2, $3, $4)
	`

	ks.db.ExecContext(ctx, query, keyID, action, performedBy, detailsJSON)
}

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

// GetAuditLog retrieves audit log entries for a key
func (ks *KeyStoreDB) GetAuditLog(ctx context.Context, keyStr string, limit int) ([]map[string]interface{}, error) {
	query := `
	SELECT al.action, al.performed_by, al.performed_at, al.details
	FROM api_key_audit_log al
	JOIN api_keys ak ON al.key_id = ak.id
	WHERE ak.key = $1
	ORDER BY al.performed_at DESC
	LIMIT $2
	`

	rows, err := ks.db.QueryContext(ctx, query, keyStr, limit)
	if err != nil {
		return nil, fmt.Errorf("failed to get audit log: %w", err)
	}
	defer rows.Close()

	var logs []map[string]interface{}
	for rows.Next() {
		var action, performedBy, details string
		var performedAt time.Time

		err := rows.Scan(&action, &performedBy, &performedAt, &details)
		if err != nil {
			return nil, err
		}

		logs = append(logs, map[string]interface{}{
			"action":       action,
			"performed_by": performedBy,
			"performed_at": performedAt,
			"details":      details,
		})
	}

	return logs, rows.Err()
}

// GetKeyByName retrieves an API key by its name
func (ks *KeyStoreDB) GetKeyByName(ctx context.Context, name string) (*APIKey, error) {
	query := `
	SELECT key, name, client_id, roles, created_at, last_used_at, expires_at, is_active, rate_limit
	FROM api_keys
	WHERE name = $1
	`

	var key APIKey
	var roles []string

	err := ks.db.QueryRowContext(ctx, query, name).Scan(
		&key.Key,
		&key.Name,
		&key.ClientID,
		pq.Array(&roles),
		&key.CreatedAt,
		&key.LastUsedAt,
		&key.ExpiresAt,
		&key.IsActive,
		&key.RateLimit,
	)

	if err == sql.ErrNoRows {
		return nil, ErrKeyNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get key by name: %w", err)
	}

	key.Roles = roles
	return &key, nil
}

// DeactivateKeyByName disables an API key by its name
func (ks *KeyStoreDB) DeactivateKeyByName(ctx context.Context, name string, deactivatedBy string) error {
	// First get the key to find its raw key for cache invalidation and audit
	key, err := ks.GetKeyByName(ctx, name)
	if err != nil {
		return err
	}

	return ks.DeactivateKey(ctx, key.Key, deactivatedBy)
}

// GetAuditLogByName retrieves audit log entries for a key identified by name
func (ks *KeyStoreDB) GetAuditLogByName(ctx context.Context, name string, limit int) ([]map[string]interface{}, error) {
	query := `
	SELECT al.action, al.performed_by, al.performed_at, al.details
	FROM api_key_audit_log al
	JOIN api_keys ak ON al.key_id = ak.id
	WHERE ak.name = $1
	ORDER BY al.performed_at DESC
	LIMIT $2
	`

	rows, err := ks.db.QueryContext(ctx, query, name, limit)
	if err != nil {
		return nil, fmt.Errorf("failed to get audit log: %w", err)
	}
	defer rows.Close()

	var logs []map[string]interface{}
	for rows.Next() {
		var action, performedBy, details string
		var performedAt time.Time

		err := rows.Scan(&action, &performedBy, &performedAt, &details)
		if err != nil {
			return nil, err
		}

		logs = append(logs, map[string]interface{}{
			"action":       action,
			"performed_by": performedBy,
			"performed_at": performedAt,
			"details":      details,
		})
	}

	return logs, rows.Err()
}

// Custom errors
var (
	ErrKeyNotFound = fmt.Errorf("API key not found")
)