git.stella-ops.org/src/Scheduler/__Libraries/StellaOps.Scheduler.Models/RunStateMachine.cs

using System;
using System.Collections.Generic;
using System.Linq;

namespace StellaOps.Scheduler.Models;

/// <summary>
/// Encapsulates allowed <see cref="RunState"/> transitions and invariants.
/// </summary>
public static class RunStateMachine
{
    private static readonly IReadOnlyDictionary<RunState, RunState[]> Adjacency = new Dictionary<RunState, RunState[]>
    {
        [RunState.Planning] = new[] { RunState.Planning, RunState.Queued, RunState.Cancelled },
        [RunState.Queued] = new[] { RunState.Queued, RunState.Running, RunState.Cancelled },
        [RunState.Running] = new[] { RunState.Running, RunState.Completed, RunState.Error, RunState.Cancelled },
        [RunState.Completed] = new[] { RunState.Completed },
        [RunState.Error] = new[] { RunState.Error },
        [RunState.Cancelled] = new[] { RunState.Cancelled },
    };

    public static bool CanTransition(RunState from, RunState to)
    {
        if (!Adjacency.TryGetValue(from, out var allowed))
        {
            return false;
        }

        return allowed.Contains(to);
    }

    public static bool IsTerminal(RunState state)
        => state is RunState.Completed or RunState.Error or RunState.Cancelled;

    /// <summary>
    /// Applies a state transition ensuring timestamps, stats, and error contracts stay consistent.
    /// </summary>
    public static Run EnsureTransition(
        Run run,
        RunState next,
        DateTimeOffset timestamp,
        Action<RunStatsBuilder>? mutateStats = null,
        string? errorMessage = null)
    {
        ArgumentNullException.ThrowIfNull(run);

        var normalizedTimestamp = Validation.NormalizeTimestamp(timestamp);
        var current = run.State;

        if (!CanTransition(current, next))
        {
            throw new InvalidOperationException($"Run state transition from '{current}' to '{next}' is not allowed.");
        }

        var statsBuilder = new RunStatsBuilder(run.Stats);
        mutateStats?.Invoke(statsBuilder);
        var newStats = statsBuilder.Build();

        var startedAt = run.StartedAt;
        var finishedAt = run.FinishedAt;

        if (current != RunState.Running && next == RunState.Running && startedAt is null)
        {
            startedAt = normalizedTimestamp;
        }

        if (IsTerminal(next))
        {
            finishedAt ??= normalizedTimestamp;
        }

        if (startedAt is { } start && start < run.CreatedAt)
        {
            throw new InvalidOperationException("Run started time cannot be earlier than created time.");
        }

        if (finishedAt is { } finish)
        {
            if (startedAt is { } startTime && finish < startTime)
            {
                throw new InvalidOperationException("Run finished time cannot be earlier than start time.");
            }

            if (!IsTerminal(next))
            {
                throw new InvalidOperationException("Finished time present but next state is not terminal.");
            }
        }

        string? nextError = null;
        if (next == RunState.Error)
        {
            var effectiveError = string.IsNullOrWhiteSpace(errorMessage) ? run.Error : errorMessage.Trim();
            if (string.IsNullOrWhiteSpace(effectiveError))
            {
                throw new InvalidOperationException("Transitioning to Error requires a non-empty error message.");
            }

            nextError = effectiveError;
        }
        else if (!string.IsNullOrWhiteSpace(errorMessage))
        {
            throw new InvalidOperationException("Error message can only be provided when transitioning to Error state.");
        }

        var updated = run with
        {
            State = next,
            Stats = newStats,
            StartedAt = startedAt,
            FinishedAt = finishedAt,
            Error = nextError,
        };

        Validate(updated);
        return updated;
    }

    public static void Validate(Run run)
    {
        ArgumentNullException.ThrowIfNull(run);

        if (run.StartedAt is { } started && started < run.CreatedAt)
        {
            throw new InvalidOperationException("Run.StartedAt cannot be earlier than CreatedAt.");
        }

        if (run.FinishedAt is { } finished)
        {
            if (run.StartedAt is { } startedAt && finished < startedAt)
            {
                throw new InvalidOperationException("Run.FinishedAt cannot be earlier than StartedAt.");
            }

            if (!IsTerminal(run.State))
            {
                throw new InvalidOperationException("Run.FinishedAt set while state is not terminal.");
            }
        }
        else if (IsTerminal(run.State))
        {
            throw new InvalidOperationException("Terminal run states must include FinishedAt.");
        }

        if (run.State == RunState.Error)
        {
            if (string.IsNullOrWhiteSpace(run.Error))
            {
                throw new InvalidOperationException("Run.Error must be populated when state is Error.");
            }
        }
        else if (!string.IsNullOrWhiteSpace(run.Error))
        {
            throw new InvalidOperationException("Run.Error must be null for non-error states.");
        }
    }
}