// Copyright (c) StellaOps. Licensed under AGPL-3.0-or-later.

using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Microsoft.Extensions.Logging;

namespace StellaOps.Attestor.Signing;

/// <summary>
/// Implementation of DSSE (Dead Simple Signing Envelope) signing service.
/// Supports ECDSA P-256, Ed25519, and RSA-PSS algorithms.
/// </summary>
public class DsseSigningService : IDsseSigningService
{
    private readonly IKeyProvider _keyProvider;
    private readonly ILogger<DsseSigningService> _logger;

    public DsseSigningService(
        IKeyProvider keyProvider,
        ILogger<DsseSigningService> logger)
    {
        _keyProvider = keyProvider ?? throw new ArgumentNullException(nameof(keyProvider));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public async Task<byte[]> SignAsync(
        byte[] payload,
        string payloadType,
        string signingKeyId,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(payload);
        ArgumentNullException.ThrowIfNull(payloadType);
        ArgumentNullException.ThrowIfNull(signingKeyId);

        _logger.LogDebug(
            "Signing payload with DSSE (type: {PayloadType}, key: {KeyId}, size: {Size} bytes)",
            payloadType, signingKeyId, payload.Length);

        try
        {
            // Step 1: Create DSSE Pre-Authentication Encoding (PAE)
            var pae = CreatePae(payloadType, payload);

            // Step 2: Sign the PAE
            var signingKey = await _keyProvider.GetSigningKeyAsync(signingKeyId, cancellationToken);
            var signature = SignPae(pae, signingKey);

            // Step 3: Build DSSE envelope
            var envelope = new DsseEnvelope(
                Payload: Convert.ToBase64String(payload),
                PayloadType: payloadType,
                Signatures: new[]
                {
                    new DsseSignature(
                        KeyId: signingKeyId,
                        Sig: Convert.ToBase64String(signature)
                    )
                }
            );

            // Step 4: Serialize envelope to JSON
            var envelopeJson = JsonSerializer.Serialize(envelope, new JsonSerializerOptions
            {
                WriteIndented = true,
                PropertyNamingPolicy = JsonNamingPolicy.CamelCase
            });

            var envelopeBytes = Encoding.UTF8.GetBytes(envelopeJson);

            _logger.LogInformation(
                "DSSE envelope created: {Size} bytes (key: {KeyId})",
                envelopeBytes.Length, signingKeyId);

            return envelopeBytes;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to sign payload with DSSE (key: {KeyId})", signingKeyId);
            throw new DsseSigningException($"DSSE signing failed for key {signingKeyId}", ex);
        }
    }

    public async Task<bool> VerifyAsync(
        byte[] dsseEnvelope,
        IReadOnlyList<string> trustedKeyIds,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(dsseEnvelope);
        ArgumentNullException.ThrowIfNull(trustedKeyIds);

        _logger.LogDebug(
            "Verifying DSSE envelope ({Size} bytes) against {Count} trusted keys",
            dsseEnvelope.Length, trustedKeyIds.Count);

        try
        {
            // Step 1: Parse DSSE envelope
            var envelopeJson = Encoding.UTF8.GetString(dsseEnvelope);
            var envelope = JsonSerializer.Deserialize<DsseEnvelope>(envelopeJson, new JsonSerializerOptions
            {
                PropertyNameCaseInsensitive = true
            });

            if (envelope == null)
            {
                _logger.LogWarning("Failed to parse DSSE envelope");
                return false;
            }

            // Step 2: Decode payload
            var payload = Convert.FromBase64String(envelope.Payload);

            // Step 3: Create PAE
            var pae = CreatePae(envelope.PayloadType, payload);

            // Step 4: Verify at least one signature matches a trusted key
            foreach (var signature in envelope.Signatures)
            {
                if (!trustedKeyIds.Contains(signature.KeyId))
                {
                    _logger.LogDebug("Skipping untrusted key: {KeyId}", signature.KeyId);
                    continue;
                }

                try
                {
                    var verificationKey = await _keyProvider.GetVerificationKeyAsync(
                        signature.KeyId,
                        cancellationToken);

                    var signatureBytes = Convert.FromBase64String(signature.Sig);
                    var isValid = VerifySignature(pae, signatureBytes, verificationKey);

                    if (isValid)
                    {
                        _logger.LogInformation(
                            "DSSE signature verified successfully (key: {KeyId})",
                            signature.KeyId);
                        return true;
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex,
                        "Failed to verify signature with key {KeyId}",
                        signature.KeyId);
                }
            }

            _logger.LogWarning("No valid signatures found in DSSE envelope");
            return false;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "DSSE verification failed");
            return false;
        }
    }

    /// <summary>
    /// Create DSSE Pre-Authentication Encoding (PAE).
    /// PAE = "DSSEv1" + SP + LEN(type) + SP + type + SP + LEN(body) + SP + body
    /// </summary>
    private byte[] CreatePae(string payloadType, byte[] payload)
    {
        using var stream = new MemoryStream();
        using var writer = new BinaryWriter(stream);

        // DSSE version prefix
        var version = Encoding.UTF8.GetBytes("DSSEv1");
        writer.Write((ulong)version.Length);
        writer.Write(version);

        // Payload type
        var typeBytes = Encoding.UTF8.GetBytes(payloadType);
        writer.Write((ulong)typeBytes.Length);
        writer.Write(typeBytes);

        // Payload body
        writer.Write((ulong)payload.Length);
        writer.Write(payload);

        return stream.ToArray();
    }

    /// <summary>
    /// Sign PAE with private key.
    /// </summary>
    private byte[] SignPae(byte[] pae, SigningKey key)
    {
        return key.Algorithm switch
        {
            SigningAlgorithm.EcdsaP256 => SignWithEcdsa(pae, key),
            SigningAlgorithm.EcdsaP384 => SignWithEcdsa(pae, key),
            SigningAlgorithm.RsaPss => SignWithRsaPss(pae, key),
            _ => throw new NotSupportedException($"Algorithm {key.Algorithm} not supported")
        };
    }

    /// <summary>
    /// Verify signature against PAE.
    /// </summary>
    private bool VerifySignature(byte[] pae, byte[] signature, VerificationKey key)
    {
        return key.Algorithm switch
        {
            SigningAlgorithm.EcdsaP256 => VerifyEcdsa(pae, signature, key),
            SigningAlgorithm.EcdsaP384 => VerifyEcdsa(pae, signature, key),
            SigningAlgorithm.RsaPss => VerifyRsaPss(pae, signature, key),
            _ => throw new NotSupportedException($"Algorithm {key.Algorithm} not supported")
        };
    }

    private byte[] SignWithEcdsa(byte[] pae, SigningKey key)
    {
        using var ecdsa = ECDsa.Create();
        ecdsa.ImportECPrivateKey(key.PrivateKeyBytes, out _);

        var hashAlgorithm = key.Algorithm == SigningAlgorithm.EcdsaP384
            ? HashAlgorithmName.SHA384
            : HashAlgorithmName.SHA256;

        return ecdsa.SignData(pae, hashAlgorithm);
    }

    private bool VerifyEcdsa(byte[] pae, byte[] signature, VerificationKey key)
    {
        using var ecdsa = ECDsa.Create();
        ecdsa.ImportSubjectPublicKeyInfo(key.PublicKeyBytes, out _);

        var hashAlgorithm = key.Algorithm == SigningAlgorithm.EcdsaP384
            ? HashAlgorithmName.SHA384
            : HashAlgorithmName.SHA256;

        return ecdsa.VerifyData(pae, signature, hashAlgorithm);
    }

    private byte[] SignWithRsaPss(byte[] pae, SigningKey key)
    {
        using var rsa = RSA.Create();
        rsa.ImportRSAPrivateKey(key.PrivateKeyBytes, out _);

        return rsa.SignData(
            pae,
            HashAlgorithmName.SHA256,
            RSASignaturePadding.Pss);
    }

    private bool VerifyRsaPss(byte[] pae, byte[] signature, VerificationKey key)
    {
        using var rsa = RSA.Create();
        rsa.ImportRSAPublicKey(key.PublicKeyBytes, out _);

        return rsa.VerifyData(
            pae,
            signature,
            HashAlgorithmName.SHA256,
            RSASignaturePadding.Pss);
    }
}

/// <summary>
/// Provides cryptographic keys for signing and verification.
/// </summary>
public interface IKeyProvider
{
    /// <summary>
    /// Get signing key (private key) for DSSE signing.
    /// </summary>
    Task<SigningKey> GetSigningKeyAsync(string keyId, CancellationToken cancellationToken = default);

    /// <summary>
    /// Get verification key (public key) for DSSE verification.
    /// </summary>
    Task<VerificationKey> GetVerificationKeyAsync(string keyId, CancellationToken cancellationToken = default);
}

/// <summary>
/// Signing key with private key material.
/// </summary>
public record SigningKey(
    string KeyId,
    SigningAlgorithm Algorithm,
    byte[] PrivateKeyBytes
);

/// <summary>
/// Verification key with public key material.
/// </summary>
public record VerificationKey(
    string KeyId,
    SigningAlgorithm Algorithm,
    byte[] PublicKeyBytes
);

/// <summary>
/// Supported signing algorithms.
/// </summary>
public enum SigningAlgorithm
{
    EcdsaP256,
    EcdsaP384,
    RsaPss
}

/// <summary>
/// Exception thrown when DSSE signing fails.
/// </summary>
public class DsseSigningException : Exception
{
    public DsseSigningException(string message) : base(message) { }
    public DsseSigningException(string message, Exception innerException) : base(message, innerException) { }
}