namespace marlin.core.models.plan;

using System.Collections.ObjectModel;
using meta;
using util;

public class ExecutionPlan
{
    public IEnumerable<ExecutionCommand> Components => [..this._commands, ..this._deferredCleanup];

    private readonly List<ExecutionCommand> _commands;
    private readonly List<ExecutionCommand> _deferredCleanup;

    public ExecutionPlan()
    {
        this._commands = new();
        this._deferredCleanup = new();
    }

    private void AddCommand(ExecutionCommand command)
    {
        if (!command.Deferred)
            this._commands.Add(command);

        else
        {
            switch (command.Role)
            {
                case CommandRole.Cleanup:
                    this._deferredCleanup.Add(command);
                    break;

                default:
                    throw new InvalidOperationException($"Cannot defer {command.Role}");
            }
        }
    }

    public static ExecutionPlan Create(Workflow workflow, Recipe recipe)
    {
        ExecutionPlan plan = new();

        AddToExecutionPlan(plan, workflow, recipe, null);

        return plan;
    }

    private static ExecutionCommand? AddToExecutionPlan(ExecutionPlan plan, Workflow workflow, Recipe recipe, ExecutionCommand? previousCommand)
    {
        foreach (var dependentRecipe in recipe.GetRequires(workflow))
            previousCommand = AddToExecutionPlan(plan, workflow, dependentRecipe, previousCommand);

        Dictionary<string, string> recipeArgs = recipe.CreateArguments();

        foreach (var step in recipe.Steps.Values)
        {
            var command = CreateCommand(step, recipe, workflow, false, recipeArgs, out var commandArgs);
            if (previousCommand != null)
            {
                command.ResultSelector = new CommandResultSelector(previousCommand,
                    previousCommand.FailBehavior switch
                    {
                        FailBehavior.Continue => CommandResult.Unconditional,
                        FailBehavior.Stop => CommandResult.Success,
                        FailBehavior.StopAfter => CommandResult.Success,

                        _ => throw new ArgumentOutOfRangeException()
                    });
            }

            plan.AddCommand(command);

            var cleanupCommand = CreateCleanupCommand(step, command, recipe, workflow, commandArgs);
            if (cleanupCommand != null)
                plan.AddCommand(cleanupCommand);

            previousCommand = command;
        }

        return previousCommand;
    }

    private static ExecutionCommand CreateCommand(
        AbstractStep step,
        Recipe recipe,
        Workflow workflow,
        bool defer,
        Dictionary<string, string> baseArgs,
        out Dictionary<string, string> createdArgs)
    {
        string? name = null;
        string? run = null;
        FailBehavior? fail = null;

        Dictionary<string, string> args = new(baseArgs);

        foreach (var stepHierarchy in step.GetAncestorsAndThis(workflow))
        {
            name ??= stepHierarchy.Name;
            run ??= stepHierarchy.Run;
            fail ??= stepHierarchy.Fail;

            foreach (var kvp in stepHierarchy.With)
                args.TryAdd(kvp.Key, kvp.Value);
        }

        name ??= step.Id;
        fail ??= FailBehavior.Stop;

        if (run == null)
            throw new Exception($"No run found for {recipe.Name}.{name}");

        run = VariableExpander.Expand(run, args);

        IReadOnlyDictionary<string, string> artifacts = step is RecipeStep rs
            ? rs.Artifacts
            : ReadOnlyDictionary<string, string>.Empty;

        ExecutionCommand stepCommand = new ExecutionCommand
        {
            Id = step.Id,
            Name = name,
            Command = run,
            Role = step switch
            {
                RecipeStep => CommandRole.Step,
                CleanupStep => CommandRole.Cleanup,
                _ => throw new ArgumentOutOfRangeException()
            },
            FailBehavior = fail.Value,
            Artifacts = artifacts,
            Deferred = defer
        };

        createdArgs = args;
        return stepCommand;
    }

    private static ExecutionCommand? CreateCleanupCommand(RecipeStep parentStep,
        ExecutionCommand parentCommand,
        Recipe recipe,
        Workflow workflow,
        Dictionary<string, string> baseArgs)
    {
        var cleanupStep = GetCleanup(parentStep, workflow);
        if (cleanupStep == null)
            return null;

        var cleanupCommand = CreateCommand(cleanupStep, recipe, workflow,
            cleanupStep.After == CleanupAfter.Recipe, baseArgs, out _);

        cleanupCommand.ResultSelector = new CommandResultSelector(parentCommand,
            parentCommand.FailBehavior switch
            {
                FailBehavior.Stop => CommandResult.Success,
                FailBehavior.StopAfter => CommandResult.Unconditional,
                FailBehavior.Continue => CommandResult.Unconditional,

                _ => throw new ArgumentOutOfRangeException()
            });

        return cleanupCommand;
    }

    private static CleanupStep? GetCleanup(RecipeStep step, Workflow workflow)
    {
        return step
            .GetAncestorsAndThis(workflow)
            .OfType<RecipeStep>()
            .Select(x => x.Cleanup)
            .FirstOrDefault(x => x != null);
    }
}