namespace Agent.Planner;

public sealed class TaskGraph(int maxTasks = 32)
{
    private readonly Dictionary<Guid, MutableTask> _tasks = [];
    private readonly List<Guid> _insertionOrder = [];
    private CriticalPathResult _criticalPath = new([], 0);

    public int MaxTasks { get; } = maxTasks > 0 ? maxTasks : throw new ArgumentOutOfRangeException(nameof(maxTasks));
    public int Count => _tasks.Count;
    public CriticalPathResult CriticalPath => _criticalPath;
    public long TotalTokenBudget => _tasks.Values.Aggregate(0L, static (total, task) => checked(total + task.BudgetClass.MaxTokens()));
    public IReadOnlyList<PlannerTask> Tasks => _insertionOrder.Select(id => _tasks[id].Snapshot()).ToArray();

    public PlannerTask AddTask(
        string description,
        TokenBudgetClass budgetClass,
        TokenEstimate estimate,
        IEnumerable<Guid>? dependencies = null,
        bool requiresExternalEffect = false,
        Guid? id = null,
        IEnumerable<string>? requiredTools = null)
    {
        if (string.IsNullOrWhiteSpace(description)) throw new PlannerValidationException("Task description is required.");
        if (_tasks.Count >= MaxTasks) throw new PlannerValidationException($"Task limit of {MaxTasks} exceeded.");
        estimate.Validate(budgetClass);
        Guid taskId = id ?? Guid.NewGuid();
        if (_tasks.ContainsKey(taskId)) throw new PlannerValidationException($"Task '{taskId}' already exists.");
        HashSet<Guid> dependencySet = dependencies?.ToHashSet() ?? [];
        HashSet<string> tools = requiredTools?.ToHashSet(StringComparer.Ordinal) ?? [];
        if (requiresExternalEffect && tools.Count == 0) throw new PlannerValidationException("An external-effect Task must name at least one required tool.");
        EnsureDependenciesExist(dependencySet);
        if (dependencySet.Contains(taskId)) throw new PlannerValidationException("A task cannot depend on itself.");

        _tasks.Add(taskId, new MutableTask(taskId, description, dependencySet, budgetClass, estimate, requiresExternalEffect, tools));
        _insertionOrder.Add(taskId);
        try
        {
            Recalculate();
        }
        catch
        {
            _tasks.Remove(taskId);
            _insertionOrder.Remove(taskId);
            throw;
        }
        return _tasks[taskId].Snapshot();
    }

    public void AddDependency(Guid taskId, Guid dependencyId)
    {
        MutableTask task = GetMutable(taskId);
        EnsureMutable(task);
        _ = GetMutable(dependencyId);
        if (!task.Dependencies.Add(dependencyId)) return;
        try
        {
            Recalculate();
        }
        catch
        {
            task.Dependencies.Remove(dependencyId);
            throw;
        }
    }

    public void SetEstimate(Guid taskId, TokenBudgetClass budgetClass, TokenEstimate estimate)
    {
        MutableTask task = GetMutable(taskId);
        EnsureMutable(task);
        estimate.Validate(budgetClass);
        (TokenBudgetClass previousBudget, TokenEstimate previousEstimate) = (task.BudgetClass, task.Estimate);
        task.BudgetClass = budgetClass;
        task.Estimate = estimate;
        try { Recalculate(); }
        catch
        {
            task.BudgetClass = previousBudget;
            task.Estimate = previousEstimate;
            throw;
        }
    }

    public void SetState(Guid taskId, PlannerTaskState state)
    {
        MutableTask task = GetMutable(taskId);
        EnsureMutable(task);
        task.State = state;
    }

    public void AddUsage(Guid taskId, long inputTokens, long outputTokens)
    {
        if (inputTokens < 0 || outputTokens < 0) throw new ArgumentOutOfRangeException(nameof(inputTokens), "Token usage cannot be negative.");
        MutableTask task = GetMutable(taskId);
        EnsureMutable(task);
        task.ActualInputTokens = checked(task.ActualInputTokens + inputTokens);
        task.ActualOutputTokens = checked(task.ActualOutputTokens + outputTokens);
    }

    public PlannerTask Get(Guid taskId) => GetMutable(taskId).Snapshot();

    private void Recalculate()
    {
        IReadOnlyList<Guid> order = TopologicalOrder();
        var costs = new Dictionary<Guid, double>();
        var predecessors = new Dictionary<Guid, Guid?>();
        foreach (Guid id in order)
        {
            MutableTask task = _tasks[id];
            Guid? best = null;
            double bestCost = 0;
            foreach (Guid dependency in task.Dependencies.OrderBy(InsertionIndex))
            {
                double candidate = costs[dependency];
                if (best is null || candidate > bestCost)
                {
                    best = dependency;
                    bestCost = candidate;
                }
            }
            costs[id] = bestCost + task.Estimate.ExpectedTokens;
            predecessors[id] = best;
        }
        if (order.Count == 0)
        {
            _criticalPath = new([], 0);
            return;
        }
        Guid end = order[0];
        foreach (Guid id in order.Skip(1)) if (costs[id] > costs[end]) end = id;
        var path = new List<Guid>();
        for (Guid? current = end; current is not null; current = predecessors[current.Value]) path.Add(current.Value);
        path.Reverse();
        _criticalPath = new(path, costs[end]);
    }

    private IReadOnlyList<Guid> TopologicalOrder()
    {
        var indegree = _tasks.ToDictionary(static pair => pair.Key, static pair => pair.Value.Dependencies.Count);
        var ready = new SortedSet<int>(_insertionOrder.Select((id, index) => (id, index)).Where(item => indegree[item.id] == 0).Select(static item => item.index));
        var result = new List<Guid>(_tasks.Count);
        while (ready.Count > 0)
        {
            int index = ready.Min;
            ready.Remove(index);
            Guid id = _insertionOrder[index];
            result.Add(id);
            foreach (Guid dependent in _insertionOrder.Where(candidate => _tasks[candidate].Dependencies.Contains(id)))
            {
                indegree[dependent]--;
                if (indegree[dependent] == 0) ready.Add(InsertionIndex(dependent));
            }
        }
        if (result.Count != _tasks.Count) throw new PlannerValidationException("TaskGraph must be a DAG; the update creates a cycle.");
        return result;
    }

    private int InsertionIndex(Guid id) => _insertionOrder.IndexOf(id);
    private MutableTask GetMutable(Guid id) => _tasks.TryGetValue(id, out MutableTask? task) ? task : throw new KeyNotFoundException($"Task '{id}' was not found.");
    private void EnsureDependenciesExist(IEnumerable<Guid> dependencies)
    {
        foreach (Guid dependency in dependencies) _ = GetMutable(dependency);
    }
    private static void EnsureMutable(MutableTask task)
    {
        if (task.State == PlannerTaskState.Completed) throw new PlannerValidationException("Completed tasks are immutable.");
    }

    private sealed class MutableTask(Guid id, string description, HashSet<Guid> dependencies, TokenBudgetClass budgetClass, TokenEstimate estimate, bool requiresExternalEffect, HashSet<string> requiredTools)
    {
        public Guid Id { get; } = id;
        public string Description { get; } = description;
        public HashSet<Guid> Dependencies { get; } = dependencies;
        public PlannerTaskState State { get; set; } = PlannerTaskState.Queued;
        public TokenBudgetClass BudgetClass { get; set; } = budgetClass;
        public TokenEstimate Estimate { get; set; } = estimate;
        public long ActualInputTokens { get; set; }
        public long ActualOutputTokens { get; set; }
        public bool RequiresExternalEffect { get; } = requiresExternalEffect;
        public HashSet<string> RequiredTools { get; } = requiredTools;
        public PlannerTask Snapshot() => new(Id, Description, Dependencies.ToHashSet(), State, BudgetClass, Estimate, ActualInputTokens, ActualOutputTokens, RequiresExternalEffect, RequiredTools.ToHashSet(StringComparer.Ordinal));
    }
}