#region License

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#endregion

using System;
using System.Collections.Generic;
using System.Security.Cryptography;
using System.Text;
using System.Threading.Tasks;
using Amazon.Runtime;
using Gremlin.Net.Driver;
using Xunit;

namespace Gremlin.Net.UnitTest.Driver
{
    public class AuthTests
    {
        private static HttpRequestContext CreateTestContext()
        {
            return new HttpRequestContext("POST", new Uri("http://localhost:8182/gremlin"),
                new Dictionary<string, string>(), new byte[] { 0x01 });
        }

        [Fact]
        public async Task BasicAuthShouldSetCorrectAuthorizationHeader()
        {
            var interceptor = Auth.Basic("user", "pass");
            var context = CreateTestContext();

            await interceptor(context);

            var expected = "Basic " + Convert.ToBase64String(Encoding.UTF8.GetBytes("user:pass"));
            Assert.Equal(expected, context.Headers["Authorization"]);
        }

        [Fact]
        public async Task BasicAuthShouldSetHeaderOnEveryInvocation()
        {
            var interceptor = Auth.Basic("user", "pass");
            var context1 = CreateTestContext();
            var context2 = CreateTestContext();

            await interceptor(context1);
            await interceptor(context2);

            var expected = "Basic " + Convert.ToBase64String(Encoding.UTF8.GetBytes("user:pass"));
            Assert.Equal(expected, context1.Headers["Authorization"]);
            Assert.Equal(expected, context2.Headers["Authorization"]);
        }

        [Fact]
        public async Task BasicAuthShouldHandleColonsInPassword()
        {
            var interceptor = Auth.Basic("user", "pass:with:colons");
            var context = CreateTestContext();

            await interceptor(context);

            var expected = "Basic " + Convert.ToBase64String(
                Encoding.UTF8.GetBytes("user:pass:with:colons"));
            Assert.Equal(expected, context.Headers["Authorization"]);
        }

        [Fact]
        public async Task BasicAuthShouldHandleUnicodeCharacters()
        {
            var interceptor = Auth.Basic("用户", "密码");
            var context = CreateTestContext();

            await interceptor(context);

            var expected = "Basic " + Convert.ToBase64String(
                Encoding.UTF8.GetBytes("用户:密码"));
            Assert.Equal(expected, context.Headers["Authorization"]);
        }

        [Fact]
        public async Task BasicAuthShouldOverwriteExistingAuthorizationHeader()
        {
            var interceptor = Auth.Basic("user", "pass");
            var context = CreateTestContext();
            context.Headers["Authorization"] = "Bearer old-token";

            await interceptor(context);

            var expected = "Basic " + Convert.ToBase64String(Encoding.UTF8.GetBytes("user:pass"));
            Assert.Equal(expected, context.Headers["Authorization"]);
        }

        [Fact]
        public async Task BasicAuthShouldHandleEmptyCredentials()
        {
            var interceptor = Auth.Basic("", "");
            var context = CreateTestContext();

            await interceptor(context);

            var expected = "Basic " + Convert.ToBase64String(Encoding.UTF8.GetBytes(":"));
            Assert.Equal(expected, context.Headers["Authorization"]);
        }

        // --- SigV4 Tests ---

        private static readonly BasicAWSCredentials TestBasicCredentials =
            new BasicAWSCredentials("MOCK_ID", "MOCK_KEY");

        private static readonly SessionAWSCredentials TestSessionCredentials =
            new SessionAWSCredentials("MOCK_ID", "MOCK_KEY", "MOCK_TOKEN");

        private static HttpRequestContext CreateSigv4TestContext(byte[]? body = null)
        {
            return new HttpRequestContext("POST", new Uri("https://example.com:8182/gremlin"),
                new Dictionary<string, string>
                {
                    { "Content-Type", "application/vnd.graphbinary-v4.0" },
                    { "Accept", "application/vnd.graphbinary-v4.0" },
                },
                body ?? new byte[] { 0x84, 0x00 });
        }

        [Fact]
        public async Task SigV4AuthShouldAddRequiredHeaders()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext();

            await interceptor(context);

            Assert.True(context.Headers.ContainsKey("Authorization"));
            Assert.True(context.Headers.ContainsKey("X-Amz-Date"));
            Assert.True(context.Headers.ContainsKey("x-amz-content-sha256"));
            Assert.True(context.Headers.ContainsKey("Host"));
        }

        [Fact]
        public async Task SigV4AuthShouldHaveCorrectAuthorizationPrefix()
        {
            var interceptor = Auth.Sigv4("gremlin-west-2", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext();

            await interceptor(context);

            Assert.StartsWith("AWS4-HMAC-SHA256 Credential=MOCK_ID", context.Headers["Authorization"]);
            Assert.Contains("gremlin-west-2/tinkerpop-sigv4/aws4_request", context.Headers["Authorization"]);
        }

        [Fact]
        public async Task SigV4AuthShouldAddSessionTokenForTemporaryCredentials()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestSessionCredentials);
            var context = CreateSigv4TestContext();

            await interceptor(context);

            Assert.True(context.Headers.ContainsKey("X-Amz-Security-Token"));
            Assert.Equal("MOCK_TOKEN", context.Headers["X-Amz-Security-Token"]);
        }

        [Fact]
        public async Task SigV4AuthShouldNotAddSessionTokenForPermanentCredentials()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext();

            await interceptor(context);

            Assert.False(context.Headers.ContainsKey("X-Amz-Security-Token"));
        }

        [Fact]
        public async Task SigV4AuthContentHashShouldMatchBodySha256()
        {
            var body = new byte[] { 0x84, 0x00, 0xFD, 0x01 };
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext(body);

            await interceptor(context);

            using var sha256 = SHA256.Create();
            var expectedHash = BitConverter.ToString(sha256.ComputeHash(body))
                .Replace("-", "").ToLowerInvariant();
            Assert.Equal(expectedHash, context.Headers["x-amz-content-sha256"]);
        }

        [Fact]
        public async Task SigV4AuthShouldHandleEmptyBody()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext(Array.Empty<byte>());

            await interceptor(context);

            Assert.True(context.Headers.ContainsKey("Authorization"));
            Assert.Equal("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
                context.Headers["x-amz-content-sha256"]);
        }

        [Fact]
        public async Task SigV4AuthShouldSetCorrectHost()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = CreateSigv4TestContext();

            await interceptor(context);

            Assert.Equal("example.com", context.Headers["Host"]);
        }

        [Fact]
        public async Task SigV4AuthShouldThrowWhenBodyIsNotByteArray()
        {
            var interceptor = Auth.Sigv4("gremlin-east-1", "tinkerpop-sigv4", TestBasicCredentials);
            var context = new HttpRequestContext("POST", new Uri("https://example.com:8182/gremlin"),
                new Dictionary<string, string>
                {
                    { "Content-Type", "application/vnd.graphbinary-v4.0" },
                },
                "not-bytes");

            var ex = await Assert.ThrowsAsync<InvalidOperationException>(
                () => interceptor(context));

            Assert.Contains("byte[]", ex.Message);
        }

        // Only host and the headers the AWS SDK adds itself are signed. Transport-managed headers
        // (accept-encoding, content-type, ...) are never signed, and the session token is signed
        // only when session credentials are used.
        private static HttpRequestContext CreateSigv4ContextWithTransportHeaders()
        {
            return new HttpRequestContext("POST", new Uri("https://example.com:8182/gremlin"),
                new Dictionary<string, string>
                {
                    { "Accept", "application/vnd.graphbinary-v4.0" },
                    { "Content-Type", "application/json" },
                    { "Accept-Encoding", "deflate" },
                    { "User-Agent", "gremlin-dotnet-test" },
                },
                Encoding.UTF8.GetBytes("{\"gremlin\":\"g.V().count()\"}"));
        }

        private static string SignedHeaders(HttpRequestContext context)
        {
            var authorization = context.Headers["Authorization"];
            const string marker = "SignedHeaders=";
            var start = authorization.IndexOf(marker, StringComparison.Ordinal) + marker.Length;
            var end = authorization.IndexOf(',', start);
            return end < 0 ? authorization.Substring(start) : authorization.Substring(start, end - start);
        }

        [Fact]
        public async Task SigV4AuthShouldSignOnlyMinimalHeaderSetForBasicCredentials()
        {
            var interceptor = Auth.Sigv4("region-1", "example-service", TestBasicCredentials);
            var context = CreateSigv4ContextWithTransportHeaders();

            await interceptor(context);

            // The .NET SDK signer includes x-amz-content-sha256 in the signed set; that is its
            // natural behavior and is intentionally left as-is.
            Assert.Equal("host;x-amz-content-sha256;x-amz-date", SignedHeaders(context));
        }

        [Fact]
        public async Task SigV4AuthShouldSignSecurityTokenForSessionCredentials()
        {
            var interceptor = Auth.Sigv4("region-1", "example-service", TestSessionCredentials);
            var context = CreateSigv4ContextWithTransportHeaders();

            await interceptor(context);

            Assert.Equal("host;x-amz-content-sha256;x-amz-date;x-amz-security-token", SignedHeaders(context));
            Assert.Equal("MOCK_TOKEN", context.Headers["X-Amz-Security-Token"]);
        }
    }
}
