using System.Collections.Concurrent;
using System.Text.Json;
using Agent.Core.Sessions;
using Agent.Planner;

namespace Agent.Core.Tools;

public interface ITaskUpdateHandler
{
    void BeginTask(Guid taskId, Guid sessionId);
    ToolExecutionResult Execute(Guid taskId, JsonElement arguments);
    void EndTask(Guid taskId);
}

public sealed class TaskGraphToolHandler(int maxTasks = 32, SessionStore? sessions = null) : ITaskUpdateHandler
{
    private readonly ConcurrentDictionary<Guid, TaskGraph> _fallbackGraphs = new();

    public void BeginTask(Guid taskId, Guid sessionId)
    {
        TaskGraph graph = new(maxTasks);
        if (sessions is not null && sessions.TryGet(sessionId, out AgentSession? session) && session is not null)
        {
            // PlanAsync may have already created a TaskGraph for this taskId.
            // If so, keep the planned graph; otherwise add a new empty one.
            session.TaskGraphs.TryAdd(taskId, graph);
        }
        else
        {
            _fallbackGraphs.TryAdd(taskId, graph);
        }
    }

    public TaskGraph? GetGraph(Guid taskId, Guid? sessionId = null)
    {
        if (sessionId is not null && sessions is not null && sessions.TryGet(sessionId.Value, out AgentSession? session) && session is not null)
            return session.TaskGraphs.TryGetValue(taskId, out TaskGraph? g) ? g : null;
        return _fallbackGraphs.TryGetValue(taskId, out TaskGraph? fallback) ? fallback : null;
    }

    public ToolExecutionResult Execute(Guid taskId, JsonElement arguments)
    {
        TaskGraph? graph = ResolveGraph(taskId) ?? throw new InvalidOperationException("TaskGraph scope is not active.");

        string operation = arguments.GetProperty("operation").GetString()!;
        switch (operation)
        {
            case "add":
                graph.AddTask(
                    RequiredString(arguments, "description"),
                    Budget(arguments),
                    Estimate(arguments),
                    Ids(arguments, "depends_on"),
                    requiresExternalEffect: OptionalBool(arguments, "requires_external_effect"),
                    id: OptionalId(arguments, "task_id"));
                break;
            case "add_dependency":
                Guid taskIdValue = RequiredId(arguments, "task_id");
                DemandStructuralMutation(graph, taskIdValue);
                Guid[] dependencies = Ids(arguments, "depends_on");
                if (dependencies.Length == 0) throw new ToolValidationException("depends_on must contain at least one task id.");
                foreach (Guid dependency in dependencies) graph.AddDependency(taskIdValue, dependency);
                break;
            case "set_estimate":
                Guid estimateTaskId = RequiredId(arguments, "task_id");
                DemandStructuralMutation(graph, estimateTaskId);
                graph.SetEstimate(estimateTaskId, Budget(arguments), Estimate(arguments));
                break;
            default:
                throw new ToolValidationException($"Unsupported task_update operation '{operation}'.");
        }

        string output = JsonSerializer.Serialize(new
        {
            tasks = graph.Tasks.Select(static task => new
            {
                task_id = task.Id,
                description = task.Description,
                depends_on = task.Dependencies,
                state = task.State.ToString().ToLowerInvariant(),
                expected_tokens = task.ExpectedTokens,
                requires_external_effect = task.RequiresExternalEffect
            }),
            critical_path = graph.CriticalPath.TaskIds,
            critical_path_tokens = graph.CriticalPath.CumulativeTokenCost,
            total_token_budget = graph.TotalTokenBudget
        });
        return new ToolExecutionResult(true, output, false);
    }

    public void EndTask(Guid taskId)
    {
        // Remove from fallback dict; session-owned graphs persist for context summary.
        _fallbackGraphs.TryRemove(taskId, out _);
    }

    private TaskGraph? ResolveGraph(Guid taskId)
    {
        // Search all sessions first, then fallback.
        if (sessions is not null)
        {
            foreach (AgentSession? session in sessions.Sessions)
            {
                if (session is not null && session.TaskGraphs.TryGetValue(taskId, out TaskGraph? g))
                    return g;
            }
        }
        return _fallbackGraphs.TryGetValue(taskId, out TaskGraph? fallback) ? fallback : null;
    }

    private static bool OptionalBool(JsonElement arguments, string property)
    {
        if (!arguments.TryGetProperty(property, out JsonElement value)) return false;
        return value.ValueKind == JsonValueKind.True;
    }

    private static void DemandStructuralMutation(TaskGraph graph, Guid taskId)
    {
        PlannerTaskState state = graph.Get(taskId).State;
        if (state is not (PlannerTaskState.Queued or PlannerTaskState.Planning))
            throw new PlannerValidationException($"TaskGraph structure or estimate cannot be changed while Task '{taskId}' is {state}.");
    }

    private static TokenBudgetClass Budget(JsonElement arguments) =>
        TokenBudgetClassExtensions.Parse(RequiredString(arguments, "budget_class"));

    private static TokenEstimate Estimate(JsonElement arguments) => new(
        RequiredInt64(arguments, "optimistic_tokens"),
        RequiredInt64(arguments, "most_likely_tokens"),
        RequiredInt64(arguments, "pessimistic_tokens"));

    private static string RequiredString(JsonElement arguments, string property) =>
        arguments.TryGetProperty(property, out JsonElement value) && value.ValueKind == JsonValueKind.String && !string.IsNullOrWhiteSpace(value.GetString())
            ? value.GetString()!
            : throw new ToolValidationException($"{property} is required for this operation.");

    private static long RequiredInt64(JsonElement arguments, string property) =>
        arguments.TryGetProperty(property, out JsonElement value) && value.ValueKind == JsonValueKind.Number
            ? value.GetInt64()
            : throw new ToolValidationException($"{property} is required for this operation.");

    private static Guid RequiredId(JsonElement arguments, string property) =>
        OptionalId(arguments, property) ?? throw new ToolValidationException($"{property} must be a UUID for this operation.");

    private static Guid? OptionalId(JsonElement arguments, string property)
    {
        if (!arguments.TryGetProperty(property, out JsonElement value) || value.ValueKind == JsonValueKind.Null) return null;
        return value.ValueKind == JsonValueKind.String && Guid.TryParse(value.GetString(), out Guid id)
            ? id
            : throw new ToolValidationException($"{property} must be a UUID.");
    }

    private static Guid[] Ids(JsonElement arguments, string property)
    {
        if (!arguments.TryGetProperty(property, out JsonElement value)) return [];
        var ids = new List<Guid>();
        foreach (JsonElement item in value.EnumerateArray())
            ids.Add(Guid.TryParse(item.GetString(), out Guid id) ? id : throw new ToolValidationException($"{property} must contain only UUIDs."));
        return [.. ids];
    }

}