using ConformanceServer.Prompts;
using ConformanceServer.Resources;
using ConformanceServer.Tools;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using ModelContextProtocol.Extensions.Tasks;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Text.Json;

namespace ModelContextProtocol.ConformanceServer;

public class Program
{
    public static async Task MainAsync(string[] args, ILoggerProvider? loggerProvider = null, CancellationToken cancellationToken = default)
    {
        var builder = WebApplication.CreateBuilder(args);

        if (loggerProvider != null)
        {
            builder.Logging.ClearProviders();
            builder.Logging.AddProvider(loggerProvider);
        }

        // Configure the default, stateful MCP server (served at "/").
        ConfigureConformanceMcpServer(builder.Services, stateless: false);

        var app = builder.Build();

        // Also expose a stateless MCP server at "/stateless" so a single conformance server can
        // serve both the legacy stateful lifecycle (at "/") and the SEP-2575 stateless lifecycle
        // (at "/stateless", which the 2026-07-28 "caching" (SEP-2549) and MRTR (SEP-2322)
        // scenarios require) from one Kestrel port.
        HandleStatelessMcp(app);

        app.MapMcp();

        app.MapGet("/health", () => "Healthy");

        await app.RunAsync(cancellationToken);
    }

    // Registers the conformance MCP server (tools, prompts, resources, filters, and handlers)
    // into the given service collection. Shared by the stateful ("/") and stateless ("/stateless")
    // servers, which expose identical behavior except that only the stateful server registers the
    // resource-subscription handlers (see below).
    private static void ConfigureConformanceMcpServer(
        IServiceCollection services,
        bool stateless)
    {
        services.AddDistributedMemoryCache();
        var mcpServerBuilder = services
            .AddMcpServer()
            .WithHttpTransport(options => options.Stateless = stateless)
            .WithDistributedCacheEventStreamStore()
            .WithTasks(
                new InMemoryMcpTaskStore
                {
                    DefaultPollIntervalMs = 50,
                    DefaultTimeToLive = TimeSpan.FromMinutes(5),
                },
                options => options.ExecutionModeSelector = request =>
                    request.Params?.Name switch
                    {
                        "slow_compute" or "protocol_error_job" or "confirm_delete" or "multi_input"
                            => McpTaskExecutionMode.Optional,
                        "failing_job"
                            => McpTaskExecutionMode.Required,
                        "test_tool_with_task" when request.Params.InputResponses is { Count: > 0 }
                            => McpTaskExecutionMode.Required,
                        _ => McpTaskExecutionMode.Synchronous,
                    })
            .WithTools<ConformanceTools>()
            .WithTools<ConformanceTaskTools>()
            .WithTools<IncompleteResultTools>()
            .WithTools<Sep2575DiagnosticTools>()
            .WithTools([ConformanceTools.CreateJsonSchema202012Tool()])
            .WithRequestFilters(filters => filters.AddCallToolFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);

                // For the test_reconnection tool, enable polling mode after the tool runs.
                // This stores the result and closes the SSE stream, so the client
                // must reconnect via GET with Last-Event-ID to retrieve the result.
                if (request.Params.Name == "test_reconnection")
                {
                    await request.EnablePollingAsync(TimeSpan.FromMilliseconds(500), cancellationToken);
                }

                return result;
            })
            // SEP-2549: advertise TTL/cacheScope caching hints on cacheable results. The
            // conformance server's tools, prompts, resources, and resource templates are the
            // same for every caller, so they are cacheable with a "public" scope.
            .AddListToolsFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);
                result.TimeToLive = TimeSpan.FromMinutes(5);
                result.CacheScope = CacheScope.Public;
                return result;
            })
            .AddListPromptsFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);
                result.TimeToLive = TimeSpan.FromMinutes(5);
                result.CacheScope = CacheScope.Public;
                return result;
            })
            .AddListResourcesFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);
                result.TimeToLive = TimeSpan.FromMinutes(5);
                result.CacheScope = CacheScope.Public;
                return result;
            })
            .AddListResourceTemplatesFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);
                result.TimeToLive = TimeSpan.FromMinutes(5);
                result.CacheScope = CacheScope.Public;
                return result;
            })
            .AddReadResourceFilter(next => async (request, cancellationToken) =>
            {
                var result = await next(request, cancellationToken);
                result.TimeToLive = TimeSpan.FromMinutes(1);
                result.CacheScope = CacheScope.Public;
                return result;
            }))
            .WithPrompts<ConformancePrompts>()
            .WithPrompts<IncompleteResultPrompts>()
            .WithResources<ConformanceResources>()
            .WithCompleteHandler(async (ctx, ct) =>
            {
                // Basic completion support - returns empty array for conformance
                // Real implementations would provide contextual suggestions
                return new CompleteResult
                {
                    Completion = new Completion
                    {
                        Values = [],
                        HasMore = false,
                        Total = 0
                    }
                };
            })
            .WithSetLoggingLevelHandler(async (ctx, ct) =>
            {
                // The SDK updates the LoggingLevel field of the McpServer
                // Send a log notification to confirm the level was set
                await ctx.Server.SendNotificationAsync("notifications/message", new LoggingMessageNotificationParams
                {
                    Level = LoggingLevel.Info,
                    Logger = "conformance-test-server",
                    Data = JsonElement.Parse($"\"Log level set to: {ctx.Params.Level}\""),
                }, cancellationToken: ct);

                return new EmptyResult();
            });

        // Resource subscriptions require a stable SessionId to key the subscription table and a
        // persistent SSE stream to deliver notifications/resources/updated, neither of which
        // exists in the stateless lifecycle. Only the stateful server registers these handlers,
        // so only it advertises the resources.subscribe capability.
        if (!stateless)
        {
            // Dictionary of session IDs to a set of resource URIs they are subscribed to. The
            // value is a ConcurrentDictionary used as a thread-safe HashSet because .NET does not
            // have a built-in concurrent HashSet.
            ConcurrentDictionary<string, ConcurrentDictionary<string, byte>> subscriptions = new();

            mcpServerBuilder
                .WithSubscribeToResourcesHandler(async (ctx, ct) =>
                {
                    if (ctx.Server.SessionId == null)
                    {
                        throw new McpException("Cannot add subscription for server with null SessionId");
                    }
                    if (ctx.Params.Uri is { } uri)
                    {
                        var sessionSubscriptions = subscriptions.GetOrAdd(ctx.Server.SessionId, _ => new());
                        sessionSubscriptions.TryAdd(uri, 0);
                    }

                    return new EmptyResult();
                })
                .WithUnsubscribeFromResourcesHandler(async (ctx, ct) =>
                {
                    if (ctx.Server.SessionId == null)
                    {
                        throw new McpException("Cannot remove subscription for server with null SessionId");
                    }
                    if (ctx.Params.Uri is { } uri)
                    {
                        subscriptions[ctx.Server.SessionId].TryRemove(uri, out _);
                    }

                    return new EmptyResult();
                });
        }
    }

    // Maps a second MCP server, configured for the stateless lifecycle, at "/stateless". It is
    // built in its own ServiceCollection so its DI (and HttpServerTransportOptions) stays isolated
    // from the stateful server registered on the main host. Adapted from
    // ModelContextProtocol.TestSseServer.Program.HandleStatelessMcp.
    private static void HandleStatelessMcp(WebApplication app)
    {
        var services = new ServiceCollection();
        services.AddLogging();
        services.AddSingleton(app.Services.GetRequiredService<ILoggerFactory>());
        services.AddSingleton(app.Services.GetRequiredService<IHostApplicationLifetime>());
        services.AddSingleton(app.Services.GetRequiredService<DiagnosticListener>());
        services.AddRoutingCore();

        ConfigureConformanceMcpServer(services, stateless: true);

        var statelessApp = new ApplicationBuilder(services.BuildServiceProvider());
        statelessApp.UseRouting();
        statelessApp.UseEndpoints(endpoints => endpoints.MapMcp("/stateless"));

        // Terminal middleware that serves "/stateless" requests the main host's routing did not
        // match. Registered before app.MapMcp() so the stateful endpoints still win for "/".
        app.Run(statelessApp.Build());
    }

    public static async Task Main(string[] args)
    {
        await MainAsync(args);
    }
}