using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Diagnostics;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json.Serialization;
using VaultSharp;
using VaultSharp.V1.SecretsEngines;
using VaultSharp.V1.SecretsEngines.Transit;
using VaultSharp.V1.Commons;
using System.Collections.Generic;

namespace Clarity.Server.Data;

public interface IPIIEntity
{
    public byte[] EncryptedDek { get; set; }
    public string Tenant { get; set; }
}

public interface IPIIVault
{
    Task<(byte[] PlaintextKey, string CiphertextKey)> GenerateTenantKEKAsync(string tenantId);
    Task<byte[]> DecryptTenantKEKAsync(string encryptedKek, string tenantId);
    Task<string> RewrapTenantKEKAsync(string encryptedKek, string tenantId);
}


public class TenantKeyProvider
{
    public byte[]? CurrentKey { get; private set; }
    public void SetKey(byte[] key) => CurrentKey = key;
}

public class PIIVault : IPIIVault
{
    private readonly IVaultClient _vaultClient;

    private const string TransitPath = "clarity-transit";
    private const string MasterKeyName = "master-key";

    public PIIVault(IVaultClient vaultClient)
    {
        _vaultClient = vaultClient;
    }

    // https://github.com/rajanadar/VaultSharp/blob/master/test/VaultSharp.Samples/Backends/Secrets/TransitSecretsBackendSamples.cs
    public async Task<(byte[] PlaintextKey, string CiphertextKey)> GenerateTenantKEKAsync(string tenantId)
    {
        var encodedContext = Convert.ToBase64String(Encoding.UTF8.GetBytes(tenantId));
        var dataKeyOptions = new DataKeyRequestOptions
        {
            DataKeyType = TransitDataKeyType.plaintext, 
            Base64EncodedContext = encodedContext
        };

        Secret<DataKeyResponse> response = await _vaultClient.V1.Secrets.Transit.GenerateDataKeyAsync(
            MasterKeyName,
            dataKeyOptions,
            TransitPath);

        byte[] rawPlaintextBytes = Convert.FromBase64String(response.Data.Base64EncodedPlainText);
        string ciphertext = response.Data.CipherText;

        return (rawPlaintextBytes, ciphertext);
    }

    public async Task<byte[]> DecryptTenantKEKAsync(string encryptedKek, string tenantId)
    {
        var encodedContext = Convert.ToBase64String(Encoding.UTF8.GetBytes(tenantId));
        var decryptOptions = new DecryptRequestOptions
        {
            CipherText = encryptedKek,
            Base64EncodedContext = encodedContext
        };

        Secret<DecryptionResponse> response = await _vaultClient.V1.Secrets.Transit.DecryptAsync(
            MasterKeyName,
            decryptOptions,
            TransitPath);

        return Convert.FromBase64String(response.Data.Base64EncodedPlainText);
    }

    public async Task<string> RewrapTenantKEKAsync(string encryptedKek, string tenantId)
    {
        var encodedContext = Convert.ToBase64String(Encoding.UTF8.GetBytes(tenantId));

        var rewrapOptions = new RewrapRequestOptions
        {
            CipherText = encryptedKek,
            Base64EncodedContext = encodedContext
        };

        var response = await _vaultClient.V1.Secrets.Transit.RewrapAsync(
            MasterKeyName,
            rewrapOptions,
            TransitPath);

        return response.Data.CipherText;
    }
}
public static class TenantSecurityExtensions
{
    public static async Task InitializeTenantSecurityAsync(this WebApplication app)
    {
        using var scope = app.Services.CreateScope();
        var db = scope.ServiceProvider.GetRequiredService<ApplicationDbContext>();
        var piiVault = scope.ServiceProvider.GetRequiredService<IPIIVault>();
        var keyProvider = scope.ServiceProvider.GetRequiredService<TenantKeyProvider>();

        var tenantId = app.Configuration["Tenant__Id"] ?? "CLARITY_01000000";

        var settings = await db.SysParams.FirstOrDefaultAsync();

        if (settings == null)
        {
            // Generate the key!
            var (plaintext, ciphertext) = await piiVault.GenerateTenantKEKAsync(tenantId);

            // Save ciphertext to DB so we have it for the next restart
            db.SysParams.Add(new SysParams { EncryptedKek = ciphertext });
            await db.SaveChangesAsync();

            // Store plaintext in RAM
            keyProvider.SetKey(plaintext);
        }
        else
        {
            // Ask Vault to unwrap the existing key using the Master Key + Context
            var plaintext = await piiVault.DecryptTenantKEKAsync(settings.EncryptedKek, tenantId);

            // Store plaintext back in RAM
            keyProvider.SetKey(plaintext);
        }
    }
}

[AttributeUsage(AttributeTargets.Property)]
public class PiiDataAttribute : Attribute
{

}


public class EnvelopeEncryptionInterceptor : SaveChangesInterceptor, IMaterializationInterceptor
{
    private readonly TenantKeyProvider _keyProvider;

    public EnvelopeEncryptionInterceptor(TenantKeyProvider keyProvider)
    {
        _keyProvider = keyProvider;
    }

    public override InterceptionResult<int> SavingChanges(DbContextEventData eventData, InterceptionResult<int> result)
    {
        EncryptEntities(eventData.Context);
        return base.SavingChanges(eventData, result);
    }

    public override ValueTask<InterceptionResult<int>> SavingChangesAsync(DbContextEventData eventData, InterceptionResult<int> result, CancellationToken cancellationToken = default)
    {
        EncryptEntities(eventData.Context);
        return base.SavingChangesAsync(eventData, result, cancellationToken);
    }

    public object InitializedInstance(MaterializationInterceptionData materializationData, object entity)
    {
        if (entity is not IPIIEntity piiEntity || _keyProvider.CurrentKey == null || piiEntity.EncryptedDek == null)
        {
            return entity;
        }

        try
        {
            // 1. Open the Envelope (Unwrap Row DEK)
            byte[] rowDek = AesGcmHelper.Decrypt(piiEntity.EncryptedDek, _keyProvider.CurrentKey);

            var props = entity.GetType().GetProperties()
                .Where(p => p.PropertyType == typeof(string) && Attribute.IsDefined(p, typeof(PiiDataAttribute)));

            foreach (var prop in props)
            {
                var cipherText = (string?)prop.GetValue(entity);
                if (!string.IsNullOrEmpty(cipherText))
                {
                    prop.SetValue(entity, AesGcmHelper.DecryptString(cipherText, rowDek));
                }
            }
        }
        catch (CryptographicException) { /* Log unauthorized tampering! */ }

        return entity;
    }

    private void EncryptEntities(DbContext? context)
    {
        if (context == null || _keyProvider.CurrentKey == null) return;

        var entries = context.ChangeTracker.Entries<IPIIEntity>()
            .Where(e => e.State == EntityState.Added || e.State == EntityState.Modified);

        foreach (var entry in entries)
        {
            var entity = entry.Entity;
            byte[] rowDek = RandomNumberGenerator.GetBytes(32);

            // 2. Wrap the Row DEK
            entity.EncryptedDek = AesGcmHelper.Encrypt(rowDek, _keyProvider.CurrentKey);

            var props = entity.GetType().GetProperties()
                .Where(p => p.PropertyType == typeof(string) && Attribute.IsDefined(p, typeof(PiiDataAttribute)));

            foreach (var prop in props)
            {
                var plainText = (string?)prop.GetValue(entity);
                if (!string.IsNullOrEmpty(plainText))
                {
                    prop.SetValue(entity, AesGcmHelper.EncryptString(plainText, rowDek));
                }
            }
        }
    }
}

/// <summary>
/// AES-GCM High-Performance Helper
/// Uses Span<byte> for speed, but returns byte[] for EF Core storage.
/// </summary>
public static class AesGcmHelper
{
    private const int NonceSize = 12; // AesGcm.NonceByteSizes.MaxSize
    private const int TagSize = 16;   // AesGcm.TagByteSizes.MaxSize

    public static byte[] Encrypt(ReadOnlySpan<byte> data, ReadOnlySpan<byte> key)
    {
        // Final buffer: [Nonce (12)] [Tag (16)] [Ciphertext (N)]
        byte[] result = new byte[NonceSize + TagSize + data.Length];

        Span<byte> nonce = result.AsSpan(0, NonceSize);
        Span<byte> tag = result.AsSpan(NonceSize, TagSize);
        Span<byte> cipher = result.AsSpan(NonceSize + TagSize);

        RandomNumberGenerator.Fill(nonce);

        using var aes = new AesGcm(key);
        aes.Encrypt(nonce, data, cipher, tag);

        return result;
    }

    public static byte[] Decrypt(ReadOnlySpan<byte> encryptedData, ReadOnlySpan<byte> key)
    {
        if (encryptedData.Length < NonceSize + TagSize)
            throw new CryptographicException("Ciphertext too short.");

        ReadOnlySpan<byte> nonce = encryptedData.Slice(0, NonceSize);
        ReadOnlySpan<byte> tag = encryptedData.Slice(NonceSize, TagSize);
        ReadOnlySpan<byte> cipher = encryptedData.Slice(NonceSize + TagSize);

        byte[] result = new byte[cipher.Length];
        using var aes = new AesGcm(key);
        aes.Decrypt(nonce, cipher, tag, result);

        return result;
    }

    public static string EncryptString(string text, byte[] key) =>
        Convert.ToBase64String(Encrypt(System.Text.Encoding.UTF8.GetBytes(text), key));

    public static string DecryptString(string cipher, byte[] key) =>
        System.Text.Encoding.UTF8.GetString(Decrypt(Convert.FromBase64String(cipher), key));
}