Redflag/aggregator-agent/internal/migration/state.go

package migration

import (
	"encoding/json"
	"fmt"
	"os"
	"time"

	"github.com/Fimeg/RedFlag/aggregator-agent/internal/config"
)

// MigrationState is imported from config package to avoid duplication

// StateManager manages migration state persistence
type StateManager struct {
	configPath string
}

// NewStateManager creates a new state manager
func NewStateManager(configPath string) *StateManager {
	return &StateManager{
		configPath: configPath,
	}
}

// LoadState loads migration state from config file
func (sm *StateManager) LoadState() (*config.MigrationState, error) {
	// Load config to get migration state
	cfg, err := sm.loadConfig()
	if err != nil {
		if os.IsNotExist(err) {
			// Fresh install - no migration state yet
			return &config.MigrationState{
				LastCompleted:      make(map[string]time.Time),
				AgentVersion:       "",
				ConfigVersion:      "",
				Timestamp:          time.Now(),
				Success:           false,
				CompletedMigrations: []string{},
			}, nil
		}
		return nil, fmt.Errorf("failed to load config: %w", err)
	}

	// Check if migration state exists in config
	if cfg.MigrationState == nil {
		return &config.MigrationState{
			LastCompleted:      make(map[string]time.Time),
			AgentVersion:       cfg.AgentVersion,
			ConfigVersion:      cfg.Version,
			Timestamp:          time.Now(),
			Success:           false,
			CompletedMigrations: []string{},
		}, nil
	}

	return cfg.MigrationState, nil
}

// SaveState saves migration state to config file
func (sm *StateManager) SaveState(state *config.MigrationState) error {
	// Load current config
	cfg, err := sm.loadConfig()
	if err != nil {
		return fmt.Errorf("failed to load config for state save: %w", err)
	}

	// Update migration state
	cfg.MigrationState = state
	state.Timestamp = time.Now()

	// Save config with updated state
	return sm.saveConfig(cfg)
}

// IsMigrationCompleted checks if a specific migration was completed
func (sm *StateManager) IsMigrationCompleted(migrationType string) (bool, error) {
	state, err := sm.LoadState()
	if err != nil {
		return false, err
	}

	// Check completed migrations list
	for _, completed := range state.CompletedMigrations {
		if completed == migrationType {
			return true, nil
		}
	}

	// Also check legacy last_completed map for backward compatibility
	if timestamp, exists := state.LastCompleted[migrationType]; exists {
		return !timestamp.IsZero(), nil
	}

	return false, nil
}

// MarkMigrationCompleted marks a migration as completed
func (sm *StateManager) MarkMigrationCompleted(migrationType string, rollbackPath string, agentVersion string) error {
	state, err := sm.LoadState()
	if err != nil {
		return err
	}

	// Update completed migrations list
	found := false
	for _, completed := range state.CompletedMigrations {
		if completed == migrationType {
			found = true
			// Update timestamp
			state.LastCompleted[migrationType] = time.Now()
			break
		}
	}

	if !found {
		state.CompletedMigrations = append(state.CompletedMigrations, migrationType)
	}

	state.LastCompleted[migrationType] = time.Now()
	state.AgentVersion = agentVersion
	state.Success = true
	if rollbackPath != "" {
		state.RollbackPath = rollbackPath
	}

	return sm.SaveState(state)
}

// CleanupOldDirectories removes old migration directories after successful migration
func (sm *StateManager) CleanupOldDirectories() error {
	oldDirs := []string{
		"/etc/aggregator",
		"/var/lib/aggregator",
	}

	for _, oldDir := range oldDirs {
		if _, err := os.Stat(oldDir); err == nil {
			fmt.Printf("[MIGRATION] Cleaning up old directory: %s\n", oldDir)
			if err := os.RemoveAll(oldDir); err != nil {
				fmt.Printf("[MIGRATION] Warning: Failed to remove old directory %s: %v\n", oldDir, err)
			}
		}
	}

	return nil
}

// loadConfig loads configuration from file
func (sm *StateManager) loadConfig() (*config.Config, error) {
	data, err := os.ReadFile(sm.configPath)
	if err != nil {
		return nil, err
	}

	var cfg config.Config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}

	return &cfg, nil
}

// saveConfig saves configuration to file
func (sm *StateManager) saveConfig(cfg *config.Config) error {
	data, err := json.MarshalIndent(cfg, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(sm.configPath, data, 0644)
}