sqrtspace-dotnet/src/SqrtSpace.SpaceTime.Core/CheckpointManager.cs

using System.Text.Json;

namespace SqrtSpace.SpaceTime.Core;

/// <summary>
/// Manages checkpointing for fault-tolerant operations
/// </summary>
public class CheckpointManager : IDisposable
{
    private readonly string _checkpointDirectory;
    private readonly CheckpointStrategy _strategy;
    private readonly int _checkpointInterval;
    private int _operationCount;
    private readonly List<string> _checkpointFiles = new();

    /// <summary>
    /// Initializes a new checkpoint manager
    /// </summary>
    /// <param name="checkpointDirectory">Directory to store checkpoints</param>
    /// <param name="strategy">Checkpointing strategy</param>
    /// <param name="totalOperations">Total expected operations (for √n calculation)</param>
    public CheckpointManager(
        string? checkpointDirectory = null,
        CheckpointStrategy strategy = CheckpointStrategy.SqrtN,
        long totalOperations = 1_000_000)
    {
        _checkpointDirectory = checkpointDirectory ?? Path.Combine(Path.GetTempPath(), $"spacetime_checkpoint_{Guid.NewGuid()}");
        _strategy = strategy;
        _checkpointInterval = SpaceTimeCalculator.CalculateCheckpointCount(totalOperations, strategy);
        
        Directory.CreateDirectory(_checkpointDirectory);
    }

    /// <summary>
    /// Checks if a checkpoint should be created
    /// </summary>
    /// <returns>True if checkpoint should be created</returns>
    public bool ShouldCheckpoint()
    {
        _operationCount++;
        
        return _strategy switch
        {
            CheckpointStrategy.None => false,
            CheckpointStrategy.SqrtN => _operationCount % _checkpointInterval == 0,
            CheckpointStrategy.Linear => _operationCount % 1000 == 0,
            CheckpointStrategy.Logarithmic => IsPowerOfTwo(_operationCount),
            _ => false
        };
    }

    /// <summary>
    /// Creates a checkpoint for the given state
    /// </summary>
    /// <typeparam name="T">Type of state to checkpoint</typeparam>
    /// <param name="state">State to save</param>
    /// <param name="checkpointId">Optional checkpoint ID</param>
    /// <returns>Path to checkpoint file</returns>
    public async Task<string> CreateCheckpointAsync<T>(T state, string? checkpointId = null)
    {
        checkpointId ??= $"checkpoint_{_operationCount}_{DateTime.UtcNow.Ticks}";
        var filePath = Path.Combine(_checkpointDirectory, $"{checkpointId}.json");
        
        var json = JsonSerializer.Serialize(state, new JsonSerializerOptions
        {
            WriteIndented = true,
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase
        });
        
        await File.WriteAllTextAsync(filePath, json);
        _checkpointFiles.Add(filePath);
        
        // Clean up old checkpoints if using √n strategy
        if (_strategy == CheckpointStrategy.SqrtN && _checkpointFiles.Count > Math.Sqrt(_operationCount))
        {
            CleanupOldCheckpoints();
        }
        
        return filePath;
    }

    /// <summary>
    /// Restores state from the latest checkpoint
    /// </summary>
    /// <typeparam name="T">Type of state to restore</typeparam>
    /// <returns>Restored state or null if no checkpoint exists</returns>
    public async Task<T?> RestoreLatestCheckpointAsync<T>()
    {
        var latestCheckpoint = Directory.GetFiles(_checkpointDirectory, "*.json")
            .OrderByDescending(f => new FileInfo(f).LastWriteTimeUtc)
            .FirstOrDefault();
            
        if (latestCheckpoint == null)
            return default;
            
        var json = await File.ReadAllTextAsync(latestCheckpoint);
        return JsonSerializer.Deserialize<T>(json, new JsonSerializerOptions
        {
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase
        });
    }

    /// <summary>
    /// Restores state from a specific checkpoint
    /// </summary>
    /// <typeparam name="T">Type of state to restore</typeparam>
    /// <param name="checkpointId">Checkpoint ID to restore</param>
    /// <returns>Restored state or null if checkpoint doesn't exist</returns>
    public async Task<T?> RestoreCheckpointAsync<T>(string checkpointId)
    {
        var filePath = Path.Combine(_checkpointDirectory, $"{checkpointId}.json");
        
        if (!File.Exists(filePath))
            return default;
            
        var json = await File.ReadAllTextAsync(filePath);
        return JsonSerializer.Deserialize<T>(json, new JsonSerializerOptions
        {
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase
        });
    }

    /// <summary>
    /// Gets the number of operations since last checkpoint
    /// </summary>
    public int OperationsSinceLastCheckpoint => _operationCount % _checkpointInterval;

    /// <summary>
    /// Saves state for a specific checkpoint and key
    /// </summary>
    /// <typeparam name="T">Type of state to save</typeparam>
    /// <param name="checkpointId">Checkpoint ID</param>
    /// <param name="key">State key</param>
    /// <param name="state">State to save</param>
    /// <param name="cancellationToken">Cancellation token</param>
    public async Task SaveStateAsync<T>(string checkpointId, string key, T state, CancellationToken cancellationToken = default) where T : class
    {
        var filePath = Path.Combine(_checkpointDirectory, $"{checkpointId}_{key}.json");
        
        var json = JsonSerializer.Serialize(state, new JsonSerializerOptions
        {
            WriteIndented = true,
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase
        });
        
        await File.WriteAllTextAsync(filePath, json, cancellationToken);
    }

    /// <summary>
    /// Loads state for a specific checkpoint and key
    /// </summary>
    /// <typeparam name="T">Type of state to load</typeparam>
    /// <param name="checkpointId">Checkpoint ID</param>
    /// <param name="key">State key</param>
    /// <param name="cancellationToken">Cancellation token</param>
    /// <returns>Loaded state or null if not found</returns>
    public async Task<T?> LoadStateAsync<T>(string checkpointId, string key, CancellationToken cancellationToken = default) where T : class
    {
        var filePath = Path.Combine(_checkpointDirectory, $"{checkpointId}_{key}.json");
        
        if (!File.Exists(filePath))
            return null;
            
        var json = await File.ReadAllTextAsync(filePath, cancellationToken);
        return JsonSerializer.Deserialize<T>(json, new JsonSerializerOptions
        {
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase
        });
    }

    /// <summary>
    /// Cleans up checkpoint files
    /// </summary>
    public void Dispose()
    {
        try
        {
            if (Directory.Exists(_checkpointDirectory))
            {
                Directory.Delete(_checkpointDirectory, recursive: true);
            }
        }
        catch
        {
            // Best effort cleanup
        }
    }

    private void CleanupOldCheckpoints()
    {
        // Keep only the most recent √n checkpoints
        var toKeep = (int)Math.Sqrt(_operationCount);
        var toDelete = _checkpointFiles
            .OrderBy(f => new FileInfo(f).LastWriteTimeUtc)
            .Take(_checkpointFiles.Count - toKeep)
            .ToList();
            
        foreach (var file in toDelete)
        {
            try
            {
                File.Delete(file);
                _checkpointFiles.Remove(file);
            }
            catch
            {
                // Best effort
            }
        }
    }

    private static bool IsPowerOfTwo(int n)
    {
        return n > 0 && (n & (n - 1)) == 0;
    }
}

/// <summary>
/// Attribute to mark methods as checkpointable
/// </summary>
[AttributeUsage(AttributeTargets.Method)]
public class CheckpointableAttribute : Attribute
{
    /// <summary>
    /// Checkpointing strategy to use
    /// </summary>
    public CheckpointStrategy Strategy { get; set; } = CheckpointStrategy.SqrtN;
    
    /// <summary>
    /// Whether to automatically restore from checkpoint on failure
    /// </summary>
    public bool AutoRestore { get; set; } = true;
    
    /// <summary>
    /// Custom checkpoint directory
    /// </summary>
    public string? CheckpointDirectory { get; set; }
}