using System.Buffers.Text;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using IdentityShroud.Core.Model;
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 class JwtService(IJwtSignerFactory signerFactory)
{

    public byte[] CreateEncodedJwt(ReadOnlySpan<byte> payloadUtf8, JwtSigAlgName algName, DecryptedSigningKey key)
    {
        // LATER might be able to improve performance using ArrayPool
        
        IJwtSigner signer = signerFactory.Create(algName);
        MemoryStream headerMemStream = new();
        Utf8JsonWriter headerWriter = new(headerMemStream);
        WriteJwtHeader(headerWriter, algName, key.Id.ToString());
        headerWriter.Flush();
        headerMemStream.Seek(0, SeekOrigin.Begin);

        int headerBase64Length = Base64Url.GetEncodedLength((int)headerMemStream.Length);
        int payloadBase64Length = Base64Url.GetEncodedLength(payloadUtf8.Length);
        var jwtData = new byte[headerBase64Length + payloadBase64Length + 1];

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

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

        int signatureBase64Length = Base64Url.GetEncodedLength(signature.Length);

        var completeJwt = new byte[jwtData.Length + 1 + signatureBase64Length];
        Array.Copy(jwtData, completeJwt, jwtData.Length);
        completeJwt[jwtData.Length] = (byte)'.';

        written = Base64Url.EncodeToUtf8(signature, completeJwt.AsSpan().Slice(jwtData.Length + 1, signatureBase64Length));
        
        if (written != signatureBase64Length)
            throw new Exception("expected signature length did not match bytes written");

        return completeJwt;
    }
    
    private static void WriteJwtHeader(Utf8JsonWriter writer, JwtSigAlgName algName, string keyId)
    {
        writer.WriteStartObject();
        writer.WriteString("typ"u8, "JWT"u8);
        writer.WriteString("alg"u8, algName.ToString());
        writer.WriteString("kid"u8, keyId);
        writer.WriteEndObject();
    }
}