using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using ModelContextProtocol;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using System.Collections.Concurrent;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace ModelContextProtocol.Extensions.Tasks;

/// <summary>
/// Extension methods for <see cref="IMcpServerBuilder"/> to enable MCP Tasks support.
/// </summary>
public static class McpTasksBuilderExtensions
{
    /// <summary>
    /// Enables MCP Tasks support backed by the specified task store.
    /// </summary>
    /// <remarks>
    /// Tasks are implemented as an alternate-result call-tool filter. Alternate-result filters registered before
    /// the Tasks filter run before task creation. Filters registered after it, along with all ordinary call-tool
    /// filters, run in the background before the tool.
    /// </remarks>
    /// <param name="builder">The server builder.</param>
    /// <param name="store">The task store.</param>
    /// <returns>The builder provided in <paramref name="builder"/>.</returns>
    public static IMcpServerBuilder WithTasks(this IMcpServerBuilder builder, IMcpTaskStore store)
        => WithTasks(builder, store, static _ => { });

    /// <summary>
    /// Enables MCP Tasks support backed by the specified task store.
    /// </summary>
    /// <param name="builder">The server builder.</param>
    /// <param name="store">The task store.</param>
    /// <param name="configure">A callback that configures per-call task execution behavior.</param>
    /// <returns>The builder provided in <paramref name="builder"/>.</returns>
    public static IMcpServerBuilder WithTasks(
        this IMcpServerBuilder builder,
        IMcpTaskStore store,
        Action<McpTasksOptions> configure)
    {
#if NET
        ArgumentNullException.ThrowIfNull(builder);
        ArgumentNullException.ThrowIfNull(store);
        ArgumentNullException.ThrowIfNull(configure);
#else
        if (builder is null) throw new ArgumentNullException(nameof(builder));
        if (store is null) throw new ArgumentNullException(nameof(store));
        if (configure is null) throw new ArgumentNullException(nameof(configure));
#endif

        McpTasksOptions taskOptions = new();
        configure(taskOptions);
#if NET
        ArgumentNullException.ThrowIfNull(taskOptions.ExecutionModeSelector);
#else
        if (taskOptions.ExecutionModeSelector is null) throw new ArgumentNullException(nameof(taskOptions.ExecutionModeSelector));
#endif

        // Resolve ILoggerFactory from the provider (rather than requiring the caller to pass one) so the
        // background task body has somewhere to report failures. It is optional: if no logging is
        // registered, the options fall back to NullLoggerFactory.
        builder.Services.AddSingleton<IConfigureOptions<McpServerOptions>>(
            sp => new McpTasksConfigureOptions(
                store,
                sp.GetRequiredService<IServiceScopeFactory>(),
                sp.GetService<ILoggerFactory>(),
                taskOptions));
        return builder;
    }

    private sealed class McpTasksConfigureOptions(
        IMcpTaskStore store,
        IServiceScopeFactory serviceScopeFactory,
        ILoggerFactory? loggerFactory,
        McpTasksOptions taskOptions) : IConfigureOptions<McpServerOptions>
    {
        private readonly IMcpTaskStore _store = store;
        private readonly IServiceScopeFactory _serviceScopeFactory = serviceScopeFactory;
        private readonly ILogger _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<McpTasksConfigureOptions>();
        private readonly McpTasksOptions _taskOptions = taskOptions;
        private readonly ConcurrentDictionary<string, CancellationTokenSource> _cancellationSources = new(StringComparer.Ordinal);

        public void Configure(McpServerOptions options)
        {
#if NET
            ArgumentNullException.ThrowIfNull(options);
#else
            if (options is null) throw new ArgumentNullException(nameof(options));
#endif

            options.Capabilities ??= new ServerCapabilities();
            options.Capabilities.Extensions ??= new Dictionary<string, object>();
            if (!options.Capabilities.Extensions.ContainsKey(TasksProtocol.ExtensionId))
            {
                options.Capabilities.Extensions[TasksProtocol.ExtensionId] = new JsonObject();
            }

            options.RequestHandlers ??= new List<McpServerRequestHandler>();
            options.RequestHandlers.Add(new McpServerRequestHandler
            {
                Method = TasksProtocol.MethodTasksGet,
                RoutingNameParameter = "taskId",
                Handler = HandleGetTask,
            });
            options.RequestHandlers.Add(new McpServerRequestHandler
            {
                Method = TasksProtocol.MethodTasksUpdate,
                RoutingNameParameter = "taskId",
                Handler = HandleUpdateTask,
            });
            options.RequestHandlers.Add(new McpServerRequestHandler
            {
                Method = TasksProtocol.MethodTasksCancel,
                RoutingNameParameter = "taskId",
                Handler = HandleCancelTask,
            });

            // Use a filter rather than a handler so it wraps around Core's tool dispatch.
            // This ensures it intercepts tool calls BEFORE the tool is invoked, allowing
            // it to spawn background execution and return the task alternate immediately.
            options.Filters.Request.CallToolWithAlternateFilters.Insert(
                options.Filters.Request.CallToolWithAlternateFilters.Count,
                async (request, next, cancellationToken) =>
                {
                    var executionMode = _taskOptions.ExecutionModeSelector(request);
                    if (executionMode is not McpTaskExecutionMode.Synchronous and
                        not McpTaskExecutionMode.Optional and
                        not McpTaskExecutionMode.Required)
                    {
                        throw new InvalidOperationException(
                            $"{nameof(McpTasksOptions.ExecutionModeSelector)} returned an invalid {nameof(McpTaskExecutionMode)} value.");
                    }

                    if (executionMode == McpTaskExecutionMode.Synchronous)
                    {
                        return await next(request, cancellationToken).ConfigureAwait(false);
                    }

                    if (!IsJuly2026OrLaterProtocolRequest(request.JsonRpcRequest) ||
                        !HasTaskExtensionOptIn(request.JsonRpcRequest))
                    {
                        if (executionMode == McpTaskExecutionMode.Required)
                        {
                            throw CreateMissingTasksCapabilityException();
                        }

                        return await next(request, cancellationToken).ConfigureAwait(false);
                    }

                    return await RunAsTaskAsync(next, request, cancellationToken).ConfigureAwait(false);
                });
        }

        private async ValueTask<ResultOrAlternate<CallToolResult>> RunAsTaskAsync(
            McpRequestHandler<CallToolRequestParams, ResultOrAlternate<CallToolResult>> next,
            RequestContext<CallToolRequestParams> request,
            CancellationToken cancellationToken)
        {
            var executionScope = _serviceScopeFactory.CreateAsyncScope();
            var executionRequest = new RequestContext<CallToolRequestParams>(
                request.Server,
                request.JsonRpcRequest,
                request.Params)
            {
                MatchedPrimitive = request.MatchedPrimitive,
                Services = executionScope.ServiceProvider,
            };

            McpTaskInfo taskInfo;
            try
            {
                taskInfo = await _store.CreateTaskAsync(cancellationToken).ConfigureAwait(false);
            }
            catch
            {
                await executionScope.DisposeAsync().ConfigureAwait(false);
                throw;
            }

            var taskId = taskInfo.TaskId;
            executionRequest.Server = request.Server.WithMcpTaskOutgoingRequestInterceptor(taskId, _store);
            var cts = new CancellationTokenSource();
            _cancellationSources[taskId] = cts;

            // Capture the token before dispatching. Cancellation can remove and dispose the source
            // before the background delegate starts.
            var taskCancellationToken = cts.Token;
            _ = Task.Run(
                () => ExecuteTaskAsync(next, executionRequest, taskId, taskCancellationToken, executionScope),
                CancellationToken.None);

            return ResultOrAlternate<CallToolResult>.FromAlternate(
                ToCreateTaskResult(taskInfo),
                McpTasksJsonContext.Default.CreateTaskResult);
        }

        private async Task ExecuteTaskAsync(
            McpRequestHandler<CallToolRequestParams, ResultOrAlternate<CallToolResult>> next,
            RequestContext<CallToolRequestParams> request,
            string taskId,
            CancellationToken taskCancellationToken,
            AsyncServiceScope executionScope)
        {
            try
            {
                try
                {
                    await ExecuteToolPipelineAsync(next, request, taskId, taskCancellationToken).ConfigureAwait(false);
                }
                finally
                {
                    await executionScope.DisposeAsync().ConfigureAwait(false);
                }
            }
            catch (Exception outer)
            {
                // Expected outcomes are recorded by ExecuteToolPipelineAsync. Reaching here means a
                // store operation, task scope, or service scope failed. Record it best-effort.
                _logger.LogError(outer, "Background execution of task '{TaskId}' terminated unexpectedly while recording its result.", taskId);

                try
                {
                    var error = new JsonRpcErrorDetail { Code = (int)McpErrorCode.InternalError, Message = outer.Message };
                    var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
                    await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
                }
                catch (Exception storeEx)
                {
                    _logger.LogError(storeEx, "Failed to record the failure of background task '{TaskId}'.", taskId);
                }
            }
            finally
            {
                if (_cancellationSources.TryRemove(taskId, out var registeredCts))
                {
                    registeredCts.Dispose();
                }
            }
        }

        private async Task ExecuteToolPipelineAsync(
            McpRequestHandler<CallToolRequestParams, ResultOrAlternate<CallToolResult>> next,
            RequestContext<CallToolRequestParams> request,
            string taskId,
            CancellationToken taskCancellationToken)
        {
            try
            {
                var augmented = await next(request, taskCancellationToken).ConfigureAwait(false);

                if (augmented.IsAlternate)
                {
                    var error = new JsonRpcErrorDetail
                    {
                        Code = (int)McpErrorCode.InternalError,
                        Message = $"{nameof(IMcpTaskStore)} is configured and the {nameof(McpServerHandlers.CallToolWithAlternateHandler)} returned IsAlternate = true. Use only one mechanism.",
                    };
                    var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
                    await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
                    return;
                }

                var resultJson = JsonSerializer.SerializeToElement(augmented.Result!, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
                await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (taskCancellationToken.IsCancellationRequested)
            {
                await _store.SetCancelledAsync(taskId, CancellationToken.None).ConfigureAwait(false);
            }
            catch (InputRequiredException)
            {
                var error = new JsonRpcErrorDetail
                {
                    Code = (int)McpErrorCode.InvalidRequest,
                    Message = "MRTR and tasks cannot be composed via [McpServerTool] yet.",
                };
                var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
                await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
            }
            catch (McpProtocolException mcpEx)
            {
                // SEP-2663 §186: protocol exceptions store as failed with JSON-RPC error shape.
                var error = new JsonRpcErrorDetail { Code = (int)mcpEx.ErrorCode, Message = mcpEx.Message };
                var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
                await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                // Non-protocol exceptions are wrapped as CallToolResult { IsError = true },
                // matching Core's BuildInitialCallToolFilter behavior.
                var errorResult = new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock
                    {
                        Text = ex is McpException
                            ? $"An error occurred invoking '{request.Params?.Name}': {ex.Message}"
                            : $"An error occurred invoking '{request.Params?.Name}'.",
                    }],
                };
                var resultJson = JsonSerializer.SerializeToElement(errorResult, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
                await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
            }
        }

        private async ValueTask<JsonNode?> HandleGetTask(JsonRpcRequest request, CancellationToken cancellationToken)
        {
            GateToJuly2026OrLaterProtocol(request, TasksProtocol.MethodTasksGet);
            GateToTasksCapability(request);

            var requestParams = request.Params?.Deserialize(McpTasksJsonContext.Default.GetTaskRequestParams)
                ?? throw new McpProtocolException("Missing params for tasks/get", McpErrorCode.InvalidParams);

            var info = await _store.GetTaskAsync(requestParams.TaskId, cancellationToken).ConfigureAwait(false);
            if (info is null)
            {
                throw new McpProtocolException($"Unknown task: '{requestParams.TaskId}'", McpErrorCode.InvalidParams);
            }

            return JsonSerializer.SerializeToNode(ToGetTaskResult(info), McpTasksJsonContext.Default.GetTaskResult);
        }

        private async ValueTask<JsonNode?> HandleUpdateTask(JsonRpcRequest request, CancellationToken cancellationToken)
        {
            GateToJuly2026OrLaterProtocol(request, TasksProtocol.MethodTasksUpdate);
            GateToTasksCapability(request);

            var taskId = request.Params?["taskId"]?.GetValue<string>()
                ?? throw new McpProtocolException("Missing params.taskId for tasks/update", McpErrorCode.InvalidParams);

            // Deserialize inputResponses using Core's options which can access the internal
            // InputResponsesCore backing property on RequestParams. The extension's source-gen
            // context cannot see that internal member.
            var inputResponses = request.Params?["inputResponses"]?.Deserialize(
                    McpJsonUtilities.DefaultOptions.GetTypeInfo<IDictionary<string, InputResponse>>())
                ?? new Dictionary<string, InputResponse>();

            await _store.ResolveInputRequestsAsync(taskId, inputResponses, cancellationToken).ConfigureAwait(false);

            return JsonSerializer.SerializeToNode(new UpdateTaskResult(), McpTasksJsonContext.Default.UpdateTaskResult);
        }

        private async ValueTask<JsonNode?> HandleCancelTask(JsonRpcRequest request, CancellationToken cancellationToken)
        {
            GateToJuly2026OrLaterProtocol(request, TasksProtocol.MethodTasksCancel);
            GateToTasksCapability(request);

            var requestParams = request.Params?.Deserialize(McpTasksJsonContext.Default.CancelTaskRequestParams)
                ?? throw new McpProtocolException("Missing params for tasks/cancel", McpErrorCode.InvalidParams);

            await _store.SetCancelledAsync(requestParams.TaskId, cancellationToken).ConfigureAwait(false);

            if (_cancellationSources.TryRemove(requestParams.TaskId, out var cts))
            {
                cts.Cancel();
                cts.Dispose();
            }

            return JsonSerializer.SerializeToNode(new CancelTaskResult(), McpTasksJsonContext.Default.CancelTaskResult);
        }

        private static void GateToJuly2026OrLaterProtocol(JsonRpcRequest request, string method)
        {
            if (!IsJuly2026OrLaterProtocolRequest(request))
            {
                throw new McpProtocolException(
                    $"The method '{method}' requires a newer protocol revision that supports tasks " +
                    $"(the '{McpProtocolVersions.July2026ProtocolVersion}' revision or later); " +
                    $"the negotiated protocol version is '{request?.Context?.ProtocolVersion ?? "(none)"}'.",
                    McpErrorCode.MethodNotFound);
            }
        }

        private static void GateToTasksCapability(JsonRpcRequest request)
        {
            if (!HasTaskExtensionOptIn(request))
            {
                throw CreateMissingTasksCapabilityException();
            }
        }

        private static MissingRequiredClientCapabilityException CreateMissingTasksCapabilityException() =>
            new(
                new ClientCapabilities
                {
                    Extensions = new Dictionary<string, object>
                    {
                        [TasksProtocol.ExtensionId] = new JsonObject(),
                    },
                },
                $"The request requires the '{TasksProtocol.ExtensionId}' client extension capability.");

        private static bool HasTaskExtensionOptIn(JsonRpcRequest request) =>
            request.Context?.ClientCapabilities?.Extensions?.ContainsKey(TasksProtocol.ExtensionId) is true ||
            request.Params?["_meta"]?[MetaKeys.ClientCapabilities]?["extensions"] is JsonObject extensions &&
                extensions.ContainsKey(TasksProtocol.ExtensionId);

        private static bool IsJuly2026OrLaterProtocolRequest(JsonRpcRequest? request) =>
            McpProtocolVersions.IsJuly2026OrLaterProtocolVersion(request?.Context?.ProtocolVersion);

        private static CreateTaskResult ToCreateTaskResult(McpTaskInfo info) => new()
        {
            TaskId = info.TaskId,
            Status = info.Status,
            CreatedAt = info.CreatedAt,
            LastUpdatedAt = info.LastUpdatedAt,
            TimeToLive = info.TimeToLive,
            PollIntervalMs = info.PollIntervalMs,
            StatusMessage = info.StatusMessage,
        };

        private static GetTaskResult ToGetTaskResult(McpTaskInfo info) => info.Status switch
        {
            McpTaskStatus.Working => new WorkingTaskResult
            {
                TaskId = info.TaskId,
                CreatedAt = info.CreatedAt,
                LastUpdatedAt = info.LastUpdatedAt,
                TimeToLive = info.TimeToLive,
                PollIntervalMs = info.PollIntervalMs,
                StatusMessage = info.StatusMessage,
                ResultType = "complete",
            },
            McpTaskStatus.Completed => new CompletedTaskResult
            {
                TaskId = info.TaskId,
                CreatedAt = info.CreatedAt,
                LastUpdatedAt = info.LastUpdatedAt,
                TimeToLive = info.TimeToLive,
                PollIntervalMs = info.PollIntervalMs,
                StatusMessage = info.StatusMessage,
                Result = info.Result ?? throw new InvalidOperationException($"Task '{info.TaskId}' is completed but has no result."),
                ResultType = "complete",
            },
            McpTaskStatus.Failed => new FailedTaskResult
            {
                TaskId = info.TaskId,
                CreatedAt = info.CreatedAt,
                LastUpdatedAt = info.LastUpdatedAt,
                TimeToLive = info.TimeToLive,
                PollIntervalMs = info.PollIntervalMs,
                StatusMessage = info.StatusMessage,
                Error = info.Error ?? throw new InvalidOperationException($"Task '{info.TaskId}' is failed but has no error."),
                ResultType = "complete",
            },
            McpTaskStatus.Cancelled => new CancelledTaskResult
            {
                TaskId = info.TaskId,
                CreatedAt = info.CreatedAt,
                LastUpdatedAt = info.LastUpdatedAt,
                TimeToLive = info.TimeToLive,
                PollIntervalMs = info.PollIntervalMs,
                StatusMessage = info.StatusMessage,
                ResultType = "complete",
            },
            McpTaskStatus.InputRequired => new InputRequiredTaskResult
            {
                TaskId = info.TaskId,
                CreatedAt = info.CreatedAt,
                LastUpdatedAt = info.LastUpdatedAt,
                TimeToLive = info.TimeToLive,
                PollIntervalMs = info.PollIntervalMs,
                StatusMessage = info.StatusMessage,
                InputRequests = info.InputRequests is IDictionary<string, InputRequest> dict
                    ? dict
                    : info.InputRequests?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value) ?? new Dictionary<string, InputRequest>(),
                ResultType = "complete",
            },
            _ => throw new InvalidOperationException($"Unknown task status: {info.Status}"),
        };
    }
}