fetch_ml/internal/auth/hybrid.go

package auth

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"
)

// HybridAuthStore combines file-based and database authentication
// Falls back to file config if database is not available
type HybridAuthStore struct {
	fileStore *Config
	dbStore   *DatabaseAuthStore
	useDB     bool
	mu        sync.RWMutex
}

// NewHybridAuthStore creates a hybrid auth store
func NewHybridAuthStore(config *Config, dbPath string) (*HybridAuthStore, error) {
	hybrid := &HybridAuthStore{
		fileStore: config,
		useDB:     false,
	}

	// Try to initialize database store
	if dbPath != "" {
		dbStore, err := NewDatabaseAuthStore(dbPath)
		if err != nil {
			log.Printf("Failed to initialize database auth store, falling back to file: %v", err)
		} else {
			hybrid.dbStore = dbStore
			hybrid.useDB = true
			log.Printf("Using database authentication store")
		}
	}

	// If database is available, migrate file-based keys to database
	if hybrid.useDB && config.Enabled && len(config.APIKeys) > 0 {
		if err := hybrid.migrateFileToDatabase(context.Background()); err != nil {
			log.Printf("Failed to migrate file keys to database: %v", err)
		}
	}
	return hybrid, nil
}

// ValidateAPIKey validates an API key using either database or file store
func (h *HybridAuthStore) ValidateAPIKey(ctx context.Context, key string) (*User, error) {
	h.mu.RLock()
	useDB := h.useDB
	h.mu.RUnlock()

	if useDB {
		user, err := h.dbStore.ValidateAPIKey(ctx, key)
		if err == nil {
			return user, nil
		}
		// Fallback to file store if database fails
	}

	// Always try file store as fallback
	if h.fileStore != nil {
		user, err := h.fileStore.ValidateAPIKey(key)
		if err == nil {
			return user, nil
		}
	}

	return nil, fmt.Errorf("invalid API key")
}

// CreateAPIKey creates an API key using the preferred store
func (h *HybridAuthStore) CreateAPIKey(
	ctx context.Context,
	userID string,
	keyHash string,
	admin bool,
	roles []string,
	permissions map[string]bool,
	expiresAt *time.Time,
) error {
	h.mu.RLock()
	useDB := h.useDB
	h.mu.RUnlock()

	if useDB {
		err := h.dbStore.CreateAPIKey(ctx, userID, keyHash, admin, roles, permissions, expiresAt)
		if err == nil {
			return nil
		}

		// If database fails, fall back to file store
		log.Printf("Database key creation failed, using file store: %v", err)
		return h.createFileAPIKey(userID, keyHash, admin, roles, permissions)
	}

	// Use file store
	return h.createFileAPIKey(userID, keyHash, admin, roles, permissions)
}

// createFileAPIKey creates an API key in the file store
func (h *HybridAuthStore) createFileAPIKey(
	userID string,
	keyHash string,
	admin bool,
	roles []string,
	permissions map[string]bool,
) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.fileStore.APIKeys == nil {
		h.fileStore.APIKeys = make(map[Username]APIKeyEntry)
	}

	h.fileStore.APIKeys[Username(userID)] = APIKeyEntry{
		Hash:        APIKeyHash(keyHash),
		Admin:       admin,
		Roles:       roles,
		Permissions: permissions,
	}

	return nil
}

// RevokeAPIKey revokes an API key
func (h *HybridAuthStore) RevokeAPIKey(ctx context.Context, userID string) error {
	h.mu.RLock()
	useDB := h.useDB
	h.mu.RUnlock()

	if useDB {
		err := h.dbStore.RevokeAPIKey(ctx, userID)
		if err == nil {
			return nil
		}

		log.Printf("Database key revocation failed: %v", err)
	}

	// Remove from file store
	h.mu.Lock()
	delete(h.fileStore.APIKeys, Username(userID))
	h.mu.Unlock()

	return nil
}

// ListUsers returns all users from the active store
func (h *HybridAuthStore) ListUsers(ctx context.Context) ([]UserInfo, error) {
	h.mu.RLock()
	useDB := h.useDB
	h.mu.RUnlock()

	if useDB {
		records, err := h.dbStore.ListUsers(ctx)
		if err == nil {
			users := make([]UserInfo, len(records))
			for i, record := range records {
				users[i] = UserInfo{
					UserID:  record.UserID,
					Admin:   record.Admin,
					KeyHash: record.KeyHash,
					Created: record.CreatedAt,
					Expires: record.ExpiresAt,
					Revoked: record.RevokedAt,
				}
			}
			return users, nil
		}

		log.Printf("Database user listing failed: %v", err)
	}

	// Use file store
	return h.listFileUsers()
}

// listFileUsers returns users from file store
func (h *HybridAuthStore) listFileUsers() ([]UserInfo, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	var users []UserInfo
	for username, entry := range h.fileStore.APIKeys {
		users = append(users, UserInfo{
			UserID:  string(username),
			Admin:   entry.Admin,
			KeyHash: string(entry.Hash),
			Created: time.Now(), // File store doesn't track creation time
		})
	}

	return users, nil
}

// migrateFileToDatabase migrates file-based keys to database
func (h *HybridAuthStore) migrateFileToDatabase(ctx context.Context) error {
	if h.fileStore == nil || len(h.fileStore.APIKeys) == 0 {
		return nil
	}

	// Use context to check for cancellation during migration
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
		// Continue with migration
	}

	log.Printf("Migrating %d API keys from file to database...", len(h.fileStore.APIKeys))

	for username, entry := range h.fileStore.APIKeys {
		userID := string(username)
		err := h.dbStore.CreateAPIKey(
			ctx,
			userID,
			string(entry.Hash),
			entry.Admin,
			entry.Roles,
			entry.Permissions,
			nil,
		)
		if err != nil {
			log.Printf("Failed to migrate key for user %s: %v", userID, err)
			continue
		}
		log.Printf("Migrated key for user: %s", userID)
	}

	log.Printf("Migration completed. Consider removing keys from config file.")
	return nil
}

// SwitchToDatabase attempts to switch to database authentication
func (h *HybridAuthStore) SwitchToDatabase(dbPath string) error {
	dbStore, err := NewDatabaseAuthStore(dbPath)
	if err != nil {
		return fmt.Errorf("failed to create database store: %w", err)
	}

	h.mu.Lock()
	defer h.mu.Unlock()

	// Close existing database if any
	if h.dbStore != nil {
		_ = h.dbStore.Close()
	}

	h.dbStore = dbStore
	h.useDB = true

	// Migrate existing keys
	if h.fileStore.Enabled && len(h.fileStore.APIKeys) > 0 {
		if err := h.migrateFileToDatabase(context.Background()); err != nil {
			log.Printf("Migration warning: %v", err)
		}
	}

	return nil
}

// Close closes the database connection
func (h *HybridAuthStore) Close() error {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.dbStore != nil {
		return h.dbStore.Close()
	}
	return nil
}

// GetDatabaseStats returns database statistics
func (h *HybridAuthStore) GetDatabaseStats(ctx context.Context) (map[string]interface{}, error) {
	h.mu.RLock()
	useDB := h.useDB
	h.mu.RUnlock()

	if !useDB {
		return map[string]interface{}{
			"store_type": "file",
			"users":      len(h.fileStore.APIKeys),
		}, nil
	}

	users, err := h.dbStore.ListUsers(ctx)
	if err != nil {
		return nil, err
	}

	return map[string]interface{}{
		"store_type": "database",
		"users":      len(users),
		"path":       "db/fetch_ml.db",
	}, nil
}