using System.Collections.Immutable;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
namespace StellaOps.Cryptography.Kms;

/// <summary>
/// File-backed KMS implementation that stores encrypted key material on disk.
/// </summary>
public sealed class FileKmsClient : IKmsClient, IDisposable
{
    private static readonly JsonSerializerOptions JsonOptions = new(JsonSerializerDefaults.Web)
    {
        WriteIndented = true,
        Converters =
        {
            new JsonStringEnumConverter(),
        },
    };
    private const int MinKeyDerivationIterations = 600_000;

    private readonly FileKmsOptions _options;
    private readonly SemaphoreSlim _mutex = new(1, 1);

    public FileKmsClient(FileKmsOptions options)
    {
        ArgumentNullException.ThrowIfNull(options);
        if (string.IsNullOrWhiteSpace(options.RootPath))
        {
            throw new ArgumentException("Root path must be provided.", nameof(options));
        }

        if (string.IsNullOrWhiteSpace(options.Password))
        {
            throw new ArgumentException("Password must be provided.", nameof(options));
        }

        _options = options;
        if (_options.KeyDerivationIterations < MinKeyDerivationIterations)
        {
            throw new ArgumentOutOfRangeException(
                nameof(options.KeyDerivationIterations),
                _options.KeyDerivationIterations,
                $"PBKDF2 iterations must be at least {MinKeyDerivationIterations:N0} to satisfy cryptographic guidance.");
        }
        Directory.CreateDirectory(_options.RootPath);
    }

    public async Task<KmsSignResult> SignAsync(
        string keyId,
        string? keyVersion,
        ReadOnlyMemory<byte> data,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);
        if (data.IsEmpty)
        {
            throw new ArgumentException("Data cannot be empty.", nameof(data));
        }

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: false).ConfigureAwait(false)
                ?? throw new InvalidOperationException($"Key '{keyId}' does not exist.");

            if (record.State == KmsKeyState.Revoked)
            {
                throw new InvalidOperationException($"Key '{keyId}' is revoked and cannot be used for signing.");
            }

            var version = ResolveVersion(record, keyVersion);
            if (version.State != KmsKeyState.Active)
            {
                throw new InvalidOperationException($"Key version '{version.VersionId}' is not active. Current state: {version.State}");
            }

            var privateKey = await LoadPrivateKeyAsync(record, version, cancellationToken).ConfigureAwait(false);
            var signature = SignData(privateKey, data.Span);
            return new KmsSignResult(record.KeyId, version.VersionId, record.Algorithm, signature);
        }
        finally
        {
            _mutex.Release();
        }
    }

    public async Task<bool> VerifyAsync(
        string keyId,
        string? keyVersion,
        ReadOnlyMemory<byte> data,
        ReadOnlyMemory<byte> signature,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);
        if (data.IsEmpty || signature.IsEmpty)
        {
            return false;
        }

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: false).ConfigureAwait(false);
            if (record is null)
            {
                return false;
            }

            var version = ResolveVersion(record, keyVersion);
            if (string.IsNullOrWhiteSpace(version.PublicKey))
            {
                return false;
            }

            return VerifyData(version.CurveName, version.PublicKey, data.Span, signature.Span);
        }
        finally
        {
            _mutex.Release();
        }
    }

    public async Task<KmsKeyMetadata> GetMetadataAsync(string keyId, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: false).ConfigureAwait(false)
                ?? throw new InvalidOperationException($"Key '{keyId}' does not exist.");
            return ToMetadata(record);
        }
        finally
        {
            _mutex.Release();
        }
    }

    public async Task<KmsKeyMaterial> ExportAsync(string keyId, string? keyVersion, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: false).ConfigureAwait(false)
                ?? throw new InvalidOperationException($"Key '{keyId}' does not exist.");

            var version = ResolveVersion(record, keyVersion);
            if (string.IsNullOrWhiteSpace(version.PublicKey))
            {
                throw new InvalidOperationException($"Key '{keyId}' version '{version.VersionId}' does not have public key material.");
            }

            var privateKey = await LoadPrivateKeyAsync(record, version, cancellationToken).ConfigureAwait(false);
            return new KmsKeyMaterial(
                record.KeyId,
                version.VersionId,
                record.Algorithm,
                version.CurveName,
                Convert.FromBase64String(privateKey.D),
                Convert.FromBase64String(privateKey.Qx),
                Convert.FromBase64String(privateKey.Qy),
                version.CreatedAt);
        }
        finally
        {
            _mutex.Release();
        }
    }

    public async Task<KmsKeyMetadata> RotateAsync(string keyId, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: true).ConfigureAwait(false)
                ?? throw new InvalidOperationException("Failed to create or load key metadata.");

            if (record.State == KmsKeyState.Revoked)
            {
                throw new InvalidOperationException($"Key '{keyId}' has been revoked and cannot be rotated.");
            }

            var timestamp = DateTimeOffset.UtcNow;
            var versionId = $"{timestamp:yyyyMMddTHHmmssfffZ}";
            var keyData = CreateKeyMaterial(record.Algorithm);

            try
            {
                var envelope = EncryptPrivateKey(keyData.PrivateBlob);
                var fileName = $"{versionId}.key.json";
                var keyPath = Path.Combine(GetKeyDirectory(keyId), fileName);
                await WriteJsonAsync(keyPath, envelope, cancellationToken).ConfigureAwait(false);

                foreach (var existing in record.Versions.Where(v => v.State == KmsKeyState.Active))
                {
                    existing.State = KmsKeyState.PendingRotation;
                }

                record.Versions.Add(new KeyVersionRecord
                {
                    VersionId = versionId,
                    State = KmsKeyState.Active,
                    CreatedAt = timestamp,
                    PublicKey = keyData.PublicKey,
                    CurveName = keyData.Curve,
                    FileName = fileName,
                });

                record.CreatedAt ??= timestamp;
                record.State = KmsKeyState.Active;
                record.ActiveVersion = versionId;

                await SaveMetadataAsync(record, cancellationToken).ConfigureAwait(false);
                return ToMetadata(record);
            }
            finally
            {
                CryptographicOperations.ZeroMemory(keyData.PrivateBlob);
            }
        }
        finally
        {
            _mutex.Release();
        }
    }

    public async Task RevokeAsync(string keyId, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        await _mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var record = await LoadOrCreateMetadataAsync(keyId, cancellationToken, createIfMissing: false).ConfigureAwait(false)
                ?? throw new InvalidOperationException($"Key '{keyId}' does not exist.");

            var timestamp = DateTimeOffset.UtcNow;
            record.State = KmsKeyState.Revoked;
            foreach (var version in record.Versions)
            {
                if (version.State != KmsKeyState.Revoked)
                {
                    version.State = KmsKeyState.Revoked;
                    version.DeactivatedAt = timestamp;
                }
            }

            await SaveMetadataAsync(record, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _mutex.Release();
        }
    }

    private static string GetMetadataPath(string root, string keyId)
        => Path.Combine(root, keyId, "metadata.json");

    private string GetKeyDirectory(string keyId)
    {
        var path = Path.Combine(_options.RootPath, keyId);
        Directory.CreateDirectory(path);
        return path;
    }

    private async Task<KeyMetadataRecord?> LoadOrCreateMetadataAsync(
        string keyId,
        CancellationToken cancellationToken,
        bool createIfMissing)
    {
        var metadataPath = GetMetadataPath(_options.RootPath, keyId);
        if (!File.Exists(metadataPath))
        {
            if (!createIfMissing)
            {
                return null;
            }

            var record = new KeyMetadataRecord
            {
                KeyId = keyId,
                Algorithm = _options.Algorithm,
                State = KmsKeyState.Active,
                CreatedAt = DateTimeOffset.UtcNow,
            };

            await SaveMetadataAsync(record, cancellationToken).ConfigureAwait(false);
            return record;
        }

        await using var stream = File.Open(metadataPath, FileMode.Open, FileAccess.Read, FileShare.Read);
        var loadedRecord = await JsonSerializer.DeserializeAsync<KeyMetadataRecord>(stream, JsonOptions, cancellationToken).ConfigureAwait(false);
        if (loadedRecord is null)
        {
            return null;
        }

        if (string.IsNullOrWhiteSpace(loadedRecord.Algorithm))
        {
            loadedRecord.Algorithm = KmsAlgorithms.Es256;
        }

        foreach (var version in loadedRecord.Versions)
        {
            if (string.IsNullOrWhiteSpace(version.CurveName))
            {
                version.CurveName = "nistP256";
            }
        }

        return loadedRecord;
    }

    private async Task SaveMetadataAsync(KeyMetadataRecord record, CancellationToken cancellationToken)
    {
        var metadataPath = GetMetadataPath(_options.RootPath, record.KeyId);
        Directory.CreateDirectory(Path.GetDirectoryName(metadataPath)!);
        await using var stream = File.Open(metadataPath, FileMode.Create, FileAccess.Write, FileShare.None);
        await JsonSerializer.SerializeAsync(stream, record, JsonOptions, cancellationToken).ConfigureAwait(false);
    }

    private async Task<EcdsaPrivateKeyRecord> LoadPrivateKeyAsync(KeyMetadataRecord record, KeyVersionRecord version, CancellationToken cancellationToken)
    {
        var keyPath = Path.Combine(GetKeyDirectory(record.KeyId), version.FileName);
        if (!File.Exists(keyPath))
        {
            throw new InvalidOperationException($"Key material for version '{version.VersionId}' was not found.");
        }

        await using var stream = File.Open(keyPath, FileMode.Open, FileAccess.Read, FileShare.Read);
        var envelope = await JsonSerializer.DeserializeAsync<KeyEnvelope>(stream, JsonOptions, cancellationToken).ConfigureAwait(false)
            ?? throw new InvalidOperationException("Key envelope could not be deserialized.");

        var payload = DecryptPrivateKey(envelope);
        try
        {
            return JsonSerializer.Deserialize<EcdsaPrivateKeyRecord>(payload, JsonOptions)
                ?? throw new InvalidOperationException("Key payload could not be deserialized.");
        }
        finally
        {
            CryptographicOperations.ZeroMemory(payload);
        }
    }

    private static KeyVersionRecord ResolveVersion(KeyMetadataRecord record, string? keyVersion)
    {
        KeyVersionRecord? version = null;
        if (!string.IsNullOrWhiteSpace(keyVersion))
        {
            version = record.Versions.SingleOrDefault(v => string.Equals(v.VersionId, keyVersion, StringComparison.Ordinal));
            if (version is null)
            {
                throw new InvalidOperationException($"Key version '{keyVersion}' does not exist for key '{record.KeyId}'.");
            }
        }
        else if (!string.IsNullOrWhiteSpace(record.ActiveVersion))
        {
            version = record.Versions.SingleOrDefault(v => string.Equals(v.VersionId, record.ActiveVersion, StringComparison.Ordinal));
        }

        version ??= record.Versions
            .Where(v => v.State == KmsKeyState.Active)
            .OrderByDescending(v => v.CreatedAt)
            .FirstOrDefault();

        if (version is null)
        {
            throw new InvalidOperationException($"Key '{record.KeyId}' does not have an active version.");
        }

        return version;
    }

    private EcdsaKeyData CreateKeyMaterial(string algorithm)
    {
        if (!string.Equals(algorithm, KmsAlgorithms.Es256, StringComparison.OrdinalIgnoreCase))
        {
            throw new NotSupportedException($"Algorithm '{algorithm}' is not supported by the file KMS driver.");
        }

        using var ecdsa = ECDsa.Create(ECCurve.NamedCurves.nistP256);
        var parameters = ecdsa.ExportParameters(true);

        var keyRecord = new EcdsaPrivateKeyRecord
        {
            Curve = "nistP256",
            D = Convert.ToBase64String(parameters.D ?? Array.Empty<byte>()),
            Qx = Convert.ToBase64String(parameters.Q.X ?? Array.Empty<byte>()),
            Qy = Convert.ToBase64String(parameters.Q.Y ?? Array.Empty<byte>()),
        };

        var privateBlob = JsonSerializer.SerializeToUtf8Bytes(keyRecord, JsonOptions);

        var qx = parameters.Q.X ?? Array.Empty<byte>();
        var qy = parameters.Q.Y ?? Array.Empty<byte>();
        var publicKey = new byte[qx.Length + qy.Length];
        Buffer.BlockCopy(qx, 0, publicKey, 0, qx.Length);
        Buffer.BlockCopy(qy, 0, publicKey, qx.Length, qy.Length);

        return new EcdsaKeyData(privateBlob, Convert.ToBase64String(publicKey), keyRecord.Curve);
    }

    private byte[] SignData(EcdsaPrivateKeyRecord privateKey, ReadOnlySpan<byte> data)
    {
        var parameters = new ECParameters
        {
            Curve = ResolveCurve(privateKey.Curve),
            D = Convert.FromBase64String(privateKey.D),
            Q = new ECPoint
            {
                X = Convert.FromBase64String(privateKey.Qx),
                Y = Convert.FromBase64String(privateKey.Qy),
            },
        };

        using var ecdsa = ECDsa.Create();
        ecdsa.ImportParameters(parameters);
        return ecdsa.SignData(data, HashAlgorithmName.SHA256);
    }

    private bool VerifyData(string curveName, string publicKeyBase64, ReadOnlySpan<byte> data, ReadOnlySpan<byte> signature)
    {
        var publicKey = Convert.FromBase64String(publicKeyBase64);
        if (publicKey.Length % 2 != 0)
        {
            return false;
        }

        var half = publicKey.Length / 2;
        var qx = publicKey[..half];
        var qy = publicKey[half..];

        var parameters = new ECParameters
        {
            Curve = ResolveCurve(curveName),
            Q = new ECPoint
            {
                X = qx,
                Y = qy,
            },
        };

        using var ecdsa = ECDsa.Create();
        ecdsa.ImportParameters(parameters);
        return ecdsa.VerifyData(data, signature, HashAlgorithmName.SHA256);
    }

    private KeyEnvelope EncryptPrivateKey(ReadOnlySpan<byte> privateKey)
    {
        var salt = RandomNumberGenerator.GetBytes(16);
        var nonce = RandomNumberGenerator.GetBytes(12);
        var key = DeriveKey(salt);

        try
        {
            var ciphertext = new byte[privateKey.Length];
            var tag = new byte[16];
            var plaintextCopy = privateKey.ToArray();

            using var aesGcm = new AesGcm(key, tag.Length);
            try
            {
                aesGcm.Encrypt(nonce, plaintextCopy, ciphertext, tag);
            }
            finally
            {
                CryptographicOperations.ZeroMemory(plaintextCopy);
            }

            return new KeyEnvelope(
                Ciphertext: Convert.ToBase64String(ciphertext),
                Nonce: Convert.ToBase64String(nonce),
                Tag: Convert.ToBase64String(tag),
                Salt: Convert.ToBase64String(salt));
        }
        finally
        {
            CryptographicOperations.ZeroMemory(key);
        }
    }

    private byte[] DecryptPrivateKey(KeyEnvelope envelope)
    {
        var salt = Convert.FromBase64String(envelope.Salt);
        var nonce = Convert.FromBase64String(envelope.Nonce);
        var tag = Convert.FromBase64String(envelope.Tag);
        var ciphertext = Convert.FromBase64String(envelope.Ciphertext);

        var key = DeriveKey(salt);
        try
        {
            var plaintext = new byte[ciphertext.Length];
            using var aesGcm = new AesGcm(key, tag.Length);
            aesGcm.Decrypt(nonce, ciphertext, tag, plaintext);

            return plaintext;
        }
        finally
        {
            CryptographicOperations.ZeroMemory(key);
        }
    }

    private byte[] DeriveKey(byte[] salt)
    {
        var key = new byte[32];
        try
        {
            var passwordBytes = Encoding.UTF8.GetBytes(_options.Password);
            try
            {
                var derived = Rfc2898DeriveBytes.Pbkdf2(passwordBytes, salt, _options.KeyDerivationIterations, HashAlgorithmName.SHA256, key.Length);
                derived.CopyTo(key.AsSpan());
                CryptographicOperations.ZeroMemory(derived);
                return key;
            }
            finally
            {
                CryptographicOperations.ZeroMemory(passwordBytes);
            }
        }
        catch
        {
            CryptographicOperations.ZeroMemory(key);
            throw;
        }
    }

    private static async Task WriteJsonAsync<T>(string path, T value, CancellationToken cancellationToken)
    {
        await using var stream = File.Open(path, FileMode.Create, FileAccess.Write, FileShare.None);
        await JsonSerializer.SerializeAsync(stream, value, JsonOptions, cancellationToken).ConfigureAwait(false);
    }

    private static KmsKeyMetadata ToMetadata(KeyMetadataRecord record)
    {
        var versions = record.Versions
            .Select(v => new KmsKeyVersionMetadata(
                v.VersionId,
                v.State,
                v.CreatedAt,
                v.DeactivatedAt,
                v.PublicKey,
                v.CurveName))
            .ToImmutableArray();

        var createdAt = record.CreatedAt ?? (versions.Length > 0 ? versions.Min(v => v.CreatedAt) : DateTimeOffset.UtcNow);
        return new KmsKeyMetadata(record.KeyId, record.Algorithm, record.State, createdAt, versions);
    }

    private sealed class KeyMetadataRecord
    {
        public string KeyId { get; set; } = string.Empty;
        public string Algorithm { get; set; } = KmsAlgorithms.Es256;
        public KmsKeyState State { get; set; } = KmsKeyState.Active;
        public DateTimeOffset? CreatedAt { get; set; }
        public string? ActiveVersion { get; set; }
        public List<KeyVersionRecord> Versions { get; set; } = new();
    }

    private sealed class KeyVersionRecord
    {
        public string VersionId { get; set; } = string.Empty;
        public KmsKeyState State { get; set; } = KmsKeyState.Active;
        public DateTimeOffset CreatedAt { get; set; }
        public DateTimeOffset? DeactivatedAt { get; set; }
        public string PublicKey { get; set; } = string.Empty;
        public string FileName { get; set; } = string.Empty;
        public string CurveName { get; set; } = string.Empty;
    }

    private sealed record KeyEnvelope(
        string Ciphertext,
        string Nonce,
        string Tag,
        string Salt);

    private sealed record EcdsaKeyData(byte[] PrivateBlob, string PublicKey, string Curve);

    private sealed class EcdsaPrivateKeyRecord
    {
        public string Curve { get; set; } = string.Empty;
        public string D { get; set; } = string.Empty;
        public string Qx { get; set; } = string.Empty;
        public string Qy { get; set; } = string.Empty;
    }

    private static ECCurve ResolveCurve(string curveName) => curveName switch
    {
        "nistP256" or "P-256" or "ES256" => ECCurve.NamedCurves.nistP256,
        _ => throw new NotSupportedException($"Curve '{curveName}' is not supported."),
    };

    public void Dispose() => _mutex.Dispose();
}