using System.Buffers.Text;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Microsoft.AspNetCore.WebUtilities;

namespace IdentityShroud.Core;

public static class JwtSignatureGenerator
{
    /// <summary>
    /// Generates a JWT signature using RS256 algorithm
    /// </summary>
    /// <param name="headerBase64Url">Base64Url encoded header</param>
    /// <param name="payloadBase64Url">Base64Url encoded payload</param>
    /// <param name="privateKey">RSA private key (PEM format or RSA parameters)</param>
    /// <returns>Base64Url encoded signature</returns>
    public static string GenerateRS256Signature(string headerBase64Url, string payloadBase64Url, RSA privateKey)
    {
        // Combine header and payload with a period
        string dataToSign = $"{headerBase64Url}.{payloadBase64Url}";

        // Convert to bytes
        byte[] dataBytes = Encoding.UTF8.GetBytes(dataToSign);

        // Sign the data using RSA-SHA256
        byte[] signatureBytes = privateKey.SignData(dataBytes, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);

        // Convert signature to Base64Url encoding
        string signature = WebEncoders.Base64UrlEncode(signatureBytes);

        return signature;
    }

    public static string GenerateCompleteJwt(string headerBase64Url, string payloadBase64Url, RSA privateKey)
    {
        string signature = GenerateRS256Signature(headerBase64Url, payloadBase64Url, privateKey);
        return $"{headerBase64Url}.{payloadBase64Url}.{signature}";
    }

}

public static class JwtCreator
{

    public static byte[] CreateEncodedJwt(ReadOnlySpan<byte> payloadUtf8, IJwtSignatureProvider signatureProvider)
    {
        MemoryStream memStream = new();
        Utf8JsonWriter writer = new(memStream);
        WriteJwtHeader(writer, signatureProvider);
        writer.Flush();
        memStream.Seek(0, SeekOrigin.Begin);

        int headerBase64Length = Base64Url.GetEncodedLength((int)memStream.Length);
        int payloadBase64Length = Base64Url.GetEncodedLength(payloadUtf8.Length);
        int signatureBase64Length = Base64Url.GetEncodedLength(signatureProvider.GetSignatureLength());
        int totalLength = headerBase64Length + 1 + payloadBase64Length + 1 + signatureBase64Length;

        var completeJwt = new byte[totalLength];

        // 
        var byteArray = new byte[memStream.Length];
        memStream.ReadExactly(byteArray, 0, (int)memStream.Length);
        int written = Base64Url.EncodeToUtf8(byteArray, completeJwt);

        if (written != headerBase64Length)
            throw new Exception("expected header length did not match bytes written");
        
        completeJwt[headerBase64Length] = (byte)'.';
        
        written = Base64Url.EncodeToUtf8(payloadUtf8, completeJwt.AsSpan().Slice(headerBase64Length + 1, payloadBase64Length));
        
        if (written != payloadBase64Length)
            throw new Exception("expected payload length did not match bytes written");
        
        completeJwt[headerBase64Length + 1 + payloadBase64Length] = (byte)'.';


        Span<byte> signature = stackalloc byte[signatureProvider.GetSignatureLength()];
        signatureProvider.CalculateSignature(
            completeJwt.AsSpan().Slice(0, headerBase64Length + 1 + payloadBase64Length),
            signature);
        
        written = Base64Url.EncodeToUtf8(signature, completeJwt.AsSpan()
            .Slice(headerBase64Length + 1 + payloadBase64Length + 1));
        
        if (written != signatureBase64Length)
            throw new Exception("expected signature length did not match bytes written");

        return completeJwt;
    }
    
    private static void WriteJwtHeader(Utf8JsonWriter writer, IJwtSignatureProvider signatureProvider)
    {
        writer.WriteStartObject();
        writer.WriteString("typ"u8, "JWT"u8);
        signatureProvider.WriteJwtHeaderFields(writer);
        writer.WriteEndObject();
    }
}