using System.Net;
using System.Text.Json.Nodes;
using ModelContextProtocol.Authentication;

namespace ModelContextProtocol.Tests;

public sealed class IdentityAssertionGrantTests : IDisposable
{
    private readonly MockHttpMessageHandler _mockHandler;
    private readonly HttpClient _httpClient;

    public IdentityAssertionGrantTests()
    {
        _mockHandler = new MockHttpMessageHandler();
        _httpClient = new HttpClient(_mockHandler);
    }

    public void Dispose()
    {
        _httpClient.Dispose();
        _mockHandler.Dispose();
    }

    #region IdentityAssertionGrantProvider Tests

    [Fact]
    public async Task IdentityAssertionGrantProvider_FullFlow_ReturnsAccessToken()
    {
        _mockHandler.Handler = request =>
        {
            var url = request.RequestUri!.ToString();

            if (url.Contains(".well-known/openid-configuration"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["issuer"] = "https://auth.mcp-server.example.com",
                    ["authorization_endpoint"] = "https://auth.mcp-server.example.com/authorize",
                    ["token_endpoint"] = "https://auth.mcp-server.example.com/token",
                    ["token_endpoint_auth_methods_supported"] = new JsonArray("client_secret_basic"),
                });
            }

            if (url.Contains("idp.example.com/token"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "mock-jag-assertion",
                    ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag",
                    ["token_type"] = "N_A",
                });
            }

            if (url.Contains("auth.mcp-server.example.com/token"))
            {
                Assert.Equal("Basic", request.Headers.Authorization?.Scheme);
                Assert.Equal(
                    Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes("mcp-client-id:mcp-client-secret")),
                    request.Headers.Authorization?.Parameter);
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "final-access-token",
                    ["token_type"] = "Bearer",
                    ["expires_in"] = 3600,
                });
            }

            return new HttpResponseMessage(HttpStatusCode.NotFound);
        };

        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "mcp-client-id",
                ClientSecret = "mcp-client-secret",
                TokenEndpointAuthMethod = "client_secret_basic",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (context, ct) =>
                {
                    Assert.Equal(new Uri("https://mcp-server.example.com"), context.ResourceUrl);
                    Assert.Equal(new Uri("https://auth.mcp-server.example.com"), context.AuthorizationServerUrl);
                    return Task.FromResult("mock-id-token");
                },
            },
            _httpClient);

        var tokens = await provider.GetAccessTokenAsync(
            resourceUrl: new Uri("https://mcp-server.example.com"),
            authorizationServerUrl: new Uri("https://auth.mcp-server.example.com"),
            TestContext.Current.CancellationToken);

        Assert.Equal("final-access-token", tokens.AccessToken);
        Assert.Equal("Bearer", tokens.TokenType);
        Assert.Equal(3600, tokens.ExpiresIn);
    }

    [Fact]
    public Task IdentityAssertionGrantProvider_DefaultsToPostRegardlessOfMetadataOrder() =>
        AssertMcpTokenEndpointAuthenticationAsync(
            new JsonArray("client_secret_basic", "client_secret_post"),
            configuredMethod: null,
            expectedAuthorizationScheme: null,
            expectSecretInBody: true);

    [Fact]
    public Task IdentityAssertionGrantProvider_FallsBackToBasicWhenPostIsUnavailable() =>
        AssertMcpTokenEndpointAuthenticationAsync(
            new JsonArray("client_secret_basic"),
            configuredMethod: null,
            expectedAuthorizationScheme: "Basic",
            expectSecretInBody: false);

    [Fact]
    public Task IdentityAssertionGrantProvider_NoneDoesNotSendClientSecret() =>
        AssertMcpTokenEndpointAuthenticationAsync(
            new JsonArray("none"),
            configuredMethod: "none",
            expectedAuthorizationScheme: null,
            expectSecretInBody: false);

    [Fact]
    public async Task IdentityAssertionGrantProvider_CachesTokens()
    {
        var mcpTokenCallCount = 0;
        _mockHandler.Handler = request =>
        {
            var url = request.RequestUri!.ToString();
            if (url.Contains(".well-known"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["authorization_endpoint"] = "https://auth.example.com/authorize",
                    ["token_endpoint"] = "https://auth.example.com/token",
                });
            }

            if (url.Contains("idp.example.com"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "mock-jag",
                    ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag",
                    ["token_type"] = "N_A",
                });
            }

            mcpTokenCallCount++;
            return JsonResponse(HttpStatusCode.OK, new JsonObject
            {
                ["access_token"] = "cached-token",
                ["token_type"] = "Bearer",
                ["expires_in"] = 3600,
            });
        };

        var idTokenCallCount = 0;
        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) =>
                {
                    idTokenCallCount++;
                    return Task.FromResult("mock-id-token");
                },
            },
            _httpClient);

        var ct = TestContext.Current.CancellationToken;

        var firstTokens = await provider.GetAccessTokenAsync(
            new Uri("https://resource.example.com"),
            new Uri("https://auth.example.com"), ct);

        var secondTokens = await provider.GetAccessTokenAsync(
            new Uri("https://resource.example.com"),
            new Uri("https://auth.example.com"), ct);

        Assert.Same(firstTokens, secondTokens);
        Assert.Equal(1, idTokenCallCount);
        Assert.Equal(1, mcpTokenCallCount);
    }

    [Fact]
    public async Task IdentityAssertionGrantProvider_InvalidateCache_ForcesRefresh()
    {
        var idTokenCallCount = 0;
        _mockHandler.Handler = request =>
        {
            var url = request.RequestUri!.ToString();
            if (url.Contains(".well-known"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["authorization_endpoint"] = "https://auth.example.com/authorize",
                    ["token_endpoint"] = "https://auth.example.com/token",
                });
            }

            if (url.Contains("idp.example.com"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "mock-jag",
                    ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag",
                    ["token_type"] = "N_A",
                });
            }

            return JsonResponse(HttpStatusCode.OK, new JsonObject
            {
                ["access_token"] = $"token-{idTokenCallCount}",
                ["token_type"] = "Bearer",
                ["expires_in"] = 3600,
            });
        };

        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) =>
                {
                    idTokenCallCount++;
                    return Task.FromResult("mock-id-token");
                },
            },
            _httpClient);

        var ct = TestContext.Current.CancellationToken;

        await provider.GetAccessTokenAsync(
            new Uri("https://resource.example.com"),
            new Uri("https://auth.example.com"), ct);

        provider.InvalidateCache();

        await provider.GetAccessTokenAsync(
            new Uri("https://resource.example.com"),
            new Uri("https://auth.example.com"), ct);

        Assert.Equal(2, idTokenCallCount);
    }

    [Fact]
    public async Task IdentityAssertionGrantProvider_IdTokenCallbackReturnsEmpty_ThrowsException()
    {
        _mockHandler.Handler = request =>
        {
            var url = request.RequestUri!.ToString();
            if (url.Contains(".well-known"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["authorization_endpoint"] = "https://auth.example.com/authorize",
                    ["token_endpoint"] = "https://auth.example.com/token",
                });
            }

            return new HttpResponseMessage(HttpStatusCode.NotFound);
        };

        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult(string.Empty),
            },
            _httpClient);

        await Assert.ThrowsAsync<IdentityAssertionGrantException>(
            () => provider.GetAccessTokenAsync(
                new Uri("https://resource.example.com"),
                new Uri("https://auth.example.com"),
                TestContext.Current.CancellationToken));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_NullOptions_ThrowsArgumentNullException()
    {
        Assert.Throws<ArgumentNullException>(() => new IdentityAssertionGrantProvider(null!, _httpClient));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_NullHttpClient_ThrowsArgumentNullException()
    {
        Assert.Throws<ArgumentNullException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult("token"),
            },
            null!));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_MissingClientId_ThrowsArgumentException()
    {
        Assert.Throws<ArgumentException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult("test"),
            },
            _httpClient));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_MissingIdTokenCallback_ThrowsArgumentNullException()
    {
        Assert.Throws<ArgumentNullException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = null!,
            },
            _httpClient));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_MissingIdpConfig_ThrowsArgumentException()
    {
        Assert.Throws<ArgumentException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpClientId = "idp-client-id",
                // Neither IdpUrl nor IdpTokenEndpoint provided
                IdTokenCallback = (_, _) => Task.FromResult("test"),
            },
            _httpClient));
    }

    [Theory]
    [InlineData("client_secret_basic")]
    [InlineData("client_secret_post")]
    public void IdentityAssertionGrantProvider_TokenEndpointAuthMethodRequiresSecret_MissingClientSecret_ThrowsArgumentException(string tokenEndpointAuthMethod)
    {
        Assert.Throws<ArgumentException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult("test"),
                TokenEndpointAuthMethod = tokenEndpointAuthMethod,
                ClientSecret = null,
            },
            _httpClient));
    }

    [Fact]
    public void IdentityAssertionGrantProvider_UnsupportedTokenEndpointAuthMethod_ThrowsArgumentException()
    {
        Assert.Throws<ArgumentException>(() => new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult("test"),
                TokenEndpointAuthMethod = "private_key_jwt",
            },
            _httpClient));
    }

    [Fact]
    public async Task IdentityAssertionGrantProvider_ConcurrentCallers_RunExchangeOnce()
    {
        // Gate the first in-flight flow so multiple callers overlap while the first holds the
        // acquisition lock. Without coalescing, each concurrent caller would run its own exchange.
        var firstEntered = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        var release = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

        var mcpTokenCallCount = 0;
        _mockHandler.AsyncHandler = async request =>
        {
            var url = request.RequestUri!.ToString();
            if (url.Contains(".well-known"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["authorization_endpoint"] = "https://auth.example.com/authorize",
                    ["token_endpoint"] = "https://auth.example.com/token",
                });
            }

            if (url.Contains("idp.example.com"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "mock-jag",
                    ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag",
                    ["token_type"] = "N_A",
                });
            }

            // MCP token endpoint: this is the exchange we expect to run exactly once.
            if (Interlocked.Increment(ref mcpTokenCallCount) == 1)
            {
                firstEntered.TrySetResult(true);
                await release.Task;
            }

            return JsonResponse(HttpStatusCode.OK, new JsonObject
            {
                ["access_token"] = "final-access-token",
                ["token_type"] = "Bearer",
                ["expires_in"] = 3600,
            });
        };

        var idTokenCallCount = 0;
        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "mcp-client-id",
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) =>
                {
                    Interlocked.Increment(ref idTokenCallCount);
                    return Task.FromResult("mock-id-token");
                },
            },
            _httpClient);

        var ct = TestContext.Current.CancellationToken;
        var resourceUrl = new Uri("https://resource.example.com");
        var authUrl = new Uri("https://auth.example.com");

        var tasks = Enumerable.Range(0, 8)
            .Select(_ => provider.GetAccessTokenAsync(resourceUrl, authUrl, ct))
            .ToArray();

        // Wait until the first flow is inside the exchange (holding the lock), then let it finish.
        var entered = await Task.WhenAny(firstEntered.Task, Task.Delay(TimeSpan.FromSeconds(30), ct));
        Assert.Same(firstEntered.Task, entered);
        release.SetResult(true);

        var results = await Task.WhenAll(tasks);

        Assert.Equal(1, mcpTokenCallCount);
        Assert.Equal(1, idTokenCallCount);
        Assert.All(results, r => Assert.Same(results[0], r));
        Assert.Equal("final-access-token", results[0].AccessToken);
    }

    #endregion

    #region IdentityAssertionGrantException Tests

    [Fact]
    public void IdentityAssertionGrantException_WithErrorCodeAndDescription_FormatsMessage()
    {
        var ex = new IdentityAssertionGrantException("Base message", "invalid_grant", "Token expired");

        Assert.Contains("Base message", ex.Message);
        Assert.Contains("invalid_grant", ex.Message);
        Assert.Contains("Token expired", ex.Message);
        Assert.Equal("invalid_grant", ex.ErrorCode);
        Assert.Equal("Token expired", ex.ErrorDescription);
    }

    [Fact]
    public void IdentityAssertionGrantException_WithErrorUri_StoresIt()
    {
        var ex = new IdentityAssertionGrantException("msg", "error", "desc", "https://docs.example.com/error");

        Assert.Equal("https://docs.example.com/error", ex.ErrorUri);
    }

    [Fact]
    public void IdentityAssertionGrantException_WithoutErrorDetails_PlainMessage()
    {
        var ex = new IdentityAssertionGrantException("Simple error");

        Assert.Equal("Simple error", ex.Message);
        Assert.Null(ex.ErrorCode);
        Assert.Null(ex.ErrorDescription);
        Assert.Null(ex.ErrorUri);
    }

    #endregion

    #region Helpers

    private async Task AssertMcpTokenEndpointAuthenticationAsync(
        JsonArray supportedMethods,
        string? configuredMethod,
        string? expectedAuthorizationScheme,
        bool expectSecretInBody)
    {
        _mockHandler.AsyncHandler = async request =>
        {
            var url = request.RequestUri!.ToString();
            if (url.Contains(".well-known"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["issuer"] = "https://auth.example.com",
                    ["authorization_endpoint"] = "https://auth.example.com/authorize",
                    ["token_endpoint"] = "https://auth.example.com/token",
                    ["token_endpoint_auth_methods_supported"] = supportedMethods,
                });
            }

            if (url.Contains("idp.example.com"))
            {
                return JsonResponse(HttpStatusCode.OK, new JsonObject
                {
                    ["access_token"] = "mock-jag",
                    ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag",
                    ["token_type"] = "N_A",
                });
            }

            Assert.Equal(expectedAuthorizationScheme, request.Headers.Authorization?.Scheme);
            var body = await request.Content!.ReadAsStringAsync(TestContext.Current.CancellationToken);
            Assert.Equal(expectSecretInBody, body.Contains("client_secret=mcp-client-secret", StringComparison.Ordinal));

            return JsonResponse(HttpStatusCode.OK, new JsonObject
            {
                ["access_token"] = "final-access-token",
                ["token_type"] = "Bearer",
            });
        };

        var provider = new IdentityAssertionGrantProvider(
            new IdentityAssertionGrantProviderOptions
            {
                ClientId = "mcp-client-id",
                ClientSecret = "mcp-client-secret",
                TokenEndpointAuthMethod = configuredMethod,
                IdpTokenEndpoint = "https://idp.example.com/token",
                IdpClientId = "idp-client-id",
                IdTokenCallback = (_, _) => Task.FromResult("mock-id-token"),
            },
            _httpClient);

        var tokens = await provider.GetAccessTokenAsync(
            new Uri("https://resource.example.com"),
            new Uri("https://auth.example.com"),
            TestContext.Current.CancellationToken);

        Assert.Equal("final-access-token", tokens.AccessToken);
    }

    private static HttpResponseMessage JsonResponse(HttpStatusCode statusCode, JsonObject payload)
    {
        return new HttpResponseMessage(statusCode)
        {
            Content = new StringContent(payload.ToJsonString(), System.Text.Encoding.UTF8, "application/json")
        };
    }

    private sealed class MockHttpMessageHandler : HttpMessageHandler
    {
        public Func<HttpRequestMessage, HttpResponseMessage>? Handler { get; set; }
        public Func<HttpRequestMessage, Task<HttpResponseMessage>>? AsyncHandler { get; set; }

        protected override async Task<HttpResponseMessage> SendAsync(
            HttpRequestMessage request, CancellationToken cancellationToken)
        {
            if (AsyncHandler is not null)
            {
                return await AsyncHandler(request);
            }

            if (Handler is not null)
            {
                return Handler(request);
            }

            return new HttpResponseMessage(HttpStatusCode.InternalServerError)
            {
                Content = new StringContent("No mock response configured")
            };
        }
    }

    #endregion
}