using Microsoft.IdentityModel.Tokens;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Crypto.Signers;
using StellaOps.Cryptography;
using System.Threading;
using System.Threading.Tasks;

namespace StellaOps.Cryptography.Plugin.PqSoft;

internal sealed class MLDsaSignerWrapper : ICryptoSigner
{
    private readonly string _keyId;
    private readonly MLDsaPrivateKeyParameters _privateKey;
    private readonly MLDsaPublicKeyParameters _publicKey;

    public MLDsaSignerWrapper(string keyId, MLDsaPrivateKeyParameters privateKey, MLDsaPublicKeyParameters publicKey)
    {
        _keyId = keyId;
        _privateKey = privateKey;
        _publicKey = publicKey;
    }

    public string KeyId => _keyId;

    public string AlgorithmId => SignatureAlgorithms.Dilithium3;

    public ValueTask<byte[]> SignAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();
        var signer = new MLDsaSigner(MLDsaParameters.ml_dsa_65, deterministic: true);
        signer.Init(true, _privateKey);
        var dataArray = data.ToArray();
        signer.BlockUpdate(dataArray, 0, dataArray.Length);
        return ValueTask.FromResult(signer.GenerateSignature());
    }

    public ValueTask<bool> VerifyAsync(ReadOnlyMemory<byte> data, ReadOnlyMemory<byte> signature, CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();
        var verifier = new MLDsaSigner(MLDsaParameters.ml_dsa_65, deterministic: true);
        verifier.Init(false, _publicKey);
        var dataArray = data.ToArray();
        verifier.BlockUpdate(dataArray, 0, dataArray.Length);
        var ok = verifier.VerifySignature(signature.ToArray());
        return ValueTask.FromResult(ok);
    }

    public JsonWebKey ExportPublicJsonWebKey()
    {
        var jwk = new JsonWebKey
        {
            Kid = _keyId,
            Alg = AlgorithmId,
            Kty = JsonWebAlgorithmsKeyTypes.Octet,
            Use = JsonWebKeyUseNames.Sig,
            Crv = "Dilithium3"
        };

        jwk.KeyOps.Add("sign");
        jwk.KeyOps.Add("verify");
        jwk.X = Base64UrlEncoder.Encode(_publicKey.GetEncoded());

        return jwk;
    }
}