using Agent.Planner;

namespace Agent.Planner.Tests;

public sealed class TaskGraphTests
{
    private static readonly TokenEstimate Small = new(60, 60, 60);

    [Fact]
    public void Known_dag_has_deterministic_critical_path()
    {
        TaskGraph graph = new();
        PlannerTask a = graph.AddTask("A", TokenBudgetClass.K1, new TokenEstimate(60, 60, 60));
        PlannerTask b = graph.AddTask("B", TokenBudgetClass.K1, new TokenEstimate(120, 120, 120), [a.Id]);
        _ = graph.AddTask("C", TokenBudgetClass.K1, new TokenEstimate(30, 30, 30), [a.Id]);
        PlannerTask d = graph.AddTask("D", TokenBudgetClass.K1, new TokenEstimate(60, 60, 60), [b.Id]);

        Assert.Equal([a.Id, b.Id, d.Id], graph.CriticalPath.TaskIds);
        Assert.Equal(240, graph.CriticalPath.CumulativeTokenCost);
    }

    [Fact]
    public void Cycle_is_rejected_without_mutating_graph()
    {
        TaskGraph graph = new();
        PlannerTask a = graph.AddTask("A", TokenBudgetClass.K1, Small);
        PlannerTask b = graph.AddTask("B", TokenBudgetClass.K1, Small, [a.Id]);

        Assert.Throws<PlannerValidationException>(() => graph.AddDependency(a.Id, b.Id));
        Assert.Empty(graph.Get(a.Id).Dependencies);
    }

    [Fact]
    public void Max_tasks_is_enforced()
    {
        TaskGraph graph = new(1);
        _ = graph.AddTask("A", TokenBudgetClass.K1, Small);
        Assert.Throws<PlannerValidationException>(() => graph.AddTask("B", TokenBudgetClass.K1, Small));
    }

    [Theory]
    [InlineData(0, 1, 1)]
    [InlineData(2, 1, 2)]
    [InlineData(1, 3, 2)]
    [InlineData(1, 2, 2048)]
    public void Invalid_estimate_is_rejected(long optimistic, long likely, long pessimistic)
    {
        TaskGraph graph = new();
        Assert.Throws<PlannerValidationException>(() => graph.AddTask("bad", TokenBudgetClass.K1, new TokenEstimate(optimistic, likely, pessimistic)));
    }

    [Fact]
    public void Completed_task_is_immutable()
    {
        TaskGraph graph = new();
        PlannerTask task = graph.AddTask("A", TokenBudgetClass.K1, Small);
        graph.SetState(task.Id, PlannerTaskState.Completed);

        Assert.Throws<PlannerValidationException>(() => graph.SetEstimate(task.Id, TokenBudgetClass.K2, new TokenEstimate(1, 2, 3)));
        Assert.Throws<PlannerValidationException>(() => graph.SetState(task.Id, PlannerTaskState.Failed));
        Assert.Throws<PlannerValidationException>(() => graph.AddUsage(task.Id, 1, 1));
    }

    [Fact]
    public void Graph_change_recalculates_cpm_and_usage_accumulates()
    {
        TaskGraph graph = new();
        PlannerTask a = graph.AddTask("A", TokenBudgetClass.K1, Small);
        PlannerTask b = graph.AddTask("B", TokenBudgetClass.K1, new TokenEstimate(30, 30, 30));
        Assert.Equal([a.Id], graph.CriticalPath.TaskIds);

        graph.AddDependency(b.Id, a.Id);
        graph.AddUsage(b.Id, 10, 5);
        graph.AddUsage(b.Id, 2, 3);

        Assert.Equal([a.Id, b.Id], graph.CriticalPath.TaskIds);
        Assert.Equal(20, graph.Get(b.Id).ActualTotalTokens);
        Assert.Equal(2_048, graph.TotalTokenBudget);
    }
}