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 *AuthConfig
	dbStore   *DatabaseAuthStore
	useDB     bool
	mu        sync.RWMutex
}

// NewHybridAuthStore creates a hybrid auth store
func NewHybridAuthStore(config *AuthConfig, 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
		}

		// If database fails, fall back to file store
		log.Printf("Database auth failed, falling back to file store: %v", err)
		return h.fileStore.ValidateAPIKey(key)
	}

	// Use file store
	return h.fileStore.ValidateAPIKey(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()
}

// UserInfo represents user information for listing
type UserInfo struct {
	UserID  string     `json:"user_id"`
	Admin   bool       `json:"admin"`
	KeyHash string     `json:"key_hash"`
	Created time.Time  `json:"created"`
	Expires *time.Time `json:"expires,omitempty"`
	Revoked *time.Time `json:"revoked,omitempty"`
}

// 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
	}

	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
}