monorepo/native/desktop/maplefile/internal/service/tokenmanager/manager.go

package tokenmanager

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"codeberg.org/mapleopentech/monorepo/cloud/maplefile-backend/pkg/maplefile/client"
	"codeberg.org/mapleopentech/monorepo/native/desktop/maplefile/internal/service/auth"
	"codeberg.org/mapleopentech/monorepo/native/desktop/maplefile/internal/usecase/session"
)

// Manager handles automatic token refresh with graceful shutdown
type Manager struct {
	config      Config
	client      *client.Client
	authService *auth.Service
	getSession  *session.GetByIdUseCase
	logger      *zap.Logger

	// Lifecycle management
	ctx       context.Context
	cancel    context.CancelFunc
	stopCh    chan struct{} // Signal to stop
	stoppedCh chan struct{} // Confirmation of stopped
	running   atomic.Bool   // Thread-safe running flag

	// Refresh state management
	mu                  sync.Mutex
	consecutiveFailures int
}

// New creates a new token manager
func New(
	config Config,
	client *client.Client,
	authService *auth.Service,
	getSession *session.GetByIdUseCase,
	logger *zap.Logger,
) *Manager {
	return &Manager{
		config:      config,
		client:      client,
		authService: authService,
		getSession:  getSession,
		logger:      logger.Named("token-manager"),
	}
}

// Start begins the token refresh background process
// Safe to call multiple times - will only start once
func (m *Manager) Start() {
	// Only start if not already running
	if !m.running.CompareAndSwap(false, true) {
		m.logger.Debug("Token manager already running, skipping start")
		return
	}

	m.ctx, m.cancel = context.WithCancel(context.Background())
	m.stopCh = make(chan struct{})
	m.stoppedCh = make(chan struct{})
	m.consecutiveFailures = 0

	m.logger.Info("Token manager starting")
	go m.refreshLoop()
}

// Stop gracefully stops the token refresh background process
// Blocks until stopped or context deadline exceeded
func (m *Manager) Stop(ctx context.Context) error {
	if !m.running.Load() {
		m.logger.Debug("Token manager not running, nothing to stop")
		return nil
	}

	m.logger.Info("Token manager stopping...")

	// Signal stop
	close(m.stopCh)

	// Wait for goroutine to finish or timeout
	select {
	case <-m.stoppedCh:
		m.logger.Info("Token manager stopped gracefully")
		return nil
	case <-ctx.Done():
		m.logger.Warn("Token manager stop timeout, forcing cancellation")
		m.cancel()
		// Wait a bit more for cancellation to take effect
		select {
		case <-m.stoppedCh:
			m.logger.Info("Token manager stopped after forced cancellation")
			return nil
		case <-time.After(100 * time.Millisecond):
			m.logger.Error("Token manager failed to stop cleanly")
			return ctx.Err()
		}
	}
}

// IsRunning returns true if the token manager is currently running
func (m *Manager) IsRunning() bool {
	return m.running.Load()
}

// refreshLoop is the background goroutine that checks and refreshes tokens
func (m *Manager) refreshLoop() {
	defer close(m.stoppedCh)
	defer m.running.Store(false)
	defer m.logger.Info("Token refresh loop exited")

	ticker := time.NewTicker(m.config.CheckInterval)
	defer ticker.Stop()

	m.logger.Info("Token refresh loop started",
		zap.Duration("check_interval", m.config.CheckInterval),
		zap.Duration("refresh_before_expiry", m.config.RefreshBeforeExpiry))

	// Do initial check immediately
	if err := m.checkAndRefresh(); err != nil {
		m.logger.Error("Initial token refresh check failed", zap.Error(err))
	}

	for {
		select {
		case <-m.stopCh:
			m.logger.Info("Token refresh loop received stop signal")
			return

		case <-m.ctx.Done():
			m.logger.Info("Token refresh loop context cancelled")
			return

		case <-ticker.C:
			if err := m.checkAndRefresh(); err != nil {
				m.logger.Error("Token refresh check failed", zap.Error(err))
			}
		}
	}
}

// checkAndRefresh checks if token refresh is needed and performs it
func (m *Manager) checkAndRefresh() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Get current session
	sess, err := m.getSession.Execute()
	if err != nil {
		m.logger.Debug("No session found, skipping refresh check", zap.Error(err))
		return nil // Not an error - user might not be logged in
	}

	if sess == nil {
		m.logger.Debug("Session is nil, skipping refresh check")
		return nil
	}

	// Check if session is still valid
	if sess.IsExpired() {
		m.logger.Warn("Session has expired, forcing logout")
		return m.forceLogout()
	}

	// Check if refresh is needed
	timeUntilExpiry := time.Until(sess.ExpiresAt)
	if timeUntilExpiry > m.config.RefreshBeforeExpiry {
		// No refresh needed yet
		if m.consecutiveFailures > 0 {
			// Reset failure counter on successful check
			m.logger.Info("Session valid, resetting failure counter")
			m.consecutiveFailures = 0
		}
		m.logger.Debug("Token refresh not needed yet",
			zap.Duration("time_until_expiry", timeUntilExpiry))
		return nil
	}

	// Refresh needed
	m.logger.Info("Token refresh needed",
		zap.Duration("time_until_expiry", timeUntilExpiry))

	// Attempt refresh (with background context, not the manager's context)
	refreshCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if err := m.client.RefreshToken(refreshCtx); err != nil {
		m.consecutiveFailures++
		m.logger.Error("Token refresh failed",
			zap.Error(err),
			zap.Int("consecutive_failures", m.consecutiveFailures),
			zap.Int("max_failures", m.config.MaxConsecutiveFailures))

		if m.consecutiveFailures >= m.config.MaxConsecutiveFailures {
			m.logger.Error("Max consecutive refresh failures reached, forcing logout")
			return m.forceLogout()
		}

		return err
	}

	// Success - reset failure counter
	m.consecutiveFailures = 0
	m.logger.Info("Token refreshed successfully",
		zap.Duration("time_until_old_expiry", timeUntilExpiry))

	return nil
}

// forceLogout forces a logout due to refresh failures
func (m *Manager) forceLogout() error {
	m.logger.Warn("Forcing logout due to token refresh issues")

	// Use background context since manager might be shutting down
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := m.authService.Logout(ctx); err != nil {
		m.logger.Error("Failed to force logout", zap.Error(err))
		return err
	}

	m.logger.Info("Force logout completed successfully")
	return nil
}