package com.github.rampi.encryption;

import com.azure.core.credential.TokenCredential;
import com.azure.identity.DefaultAzureCredentialBuilder;
import com.azure.security.keyvault.keys.KeyClient;
import com.azure.security.keyvault.keys.KeyClientBuilder;
import com.azure.security.keyvault.keys.cryptography.CryptographyClient;
import com.azure.security.keyvault.keys.cryptography.CryptographyClientBuilder;
import com.azure.security.keyvault.keys.cryptography.models.EncryptResult;
import com.azure.security.keyvault.keys.cryptography.models.EncryptionAlgorithm;
import com.azure.security.keyvault.keys.models.KeyVaultKey;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.Base64;

public class App {

    public static void main(String[] args) {
        // Key vault details
        String keyVaultName = "my-keyvault";
        String keyVaultUri = "https://" + keyVaultName + ".vault.azure.net/";
        String keyName = "my-asymmetric-key";

        try {
            // Create a credential using DefaultAzureCredential
            // This will use environment variables, managed identity, or developer credentials
            TokenCredential credential = new DefaultAzureCredentialBuilder().build();

            // Create a KeyClient to interact with Key Vault
            KeyClient keyClient = new KeyClientBuilder()
                    .vaultUrl(keyVaultUri)
                    .credential(credential)
                    .buildClient();

            // Get the key from Key Vault
            KeyVaultKey key = keyClient.getKey(keyName);

            System.out.println("Successfully retrieved key: " + key.getName());
            System.out.println("Key type: " + key.getKeyType());
            System.out.println("Key ID: " + key.getId());

            // Generate a new symmetric key with Google Tink AES-256-GCM
            // Initialize Tink
            com.google.crypto.tink.config.TinkConfig.register();

            // Generate a new AES-GCM key
            com.google.crypto.tink.KeysetHandle keysetHandle = com.google.crypto.tink.KeysetHandle.generateNew(
                    com.google.crypto.tink.aead.AesGcmKeyManager.aes256GcmTemplate());

            // Get the AEAD primitive from the keyset
            com.google.crypto.tink.Aead aead = keysetHandle.getPrimitive(com.google.crypto.tink.Aead.class);

            System.out.println("Successfully generated new symmetric AES-256-GCM key");

            // export the keyset as JSON string
            String keysetString;
            java.io.ByteArrayOutputStream outputStream = new java.io.ByteArrayOutputStream();
            com.google.crypto.tink.CleartextKeysetHandle.write(
                    keysetHandle,
                    com.google.crypto.tink.JsonKeysetWriter.withOutputStream(outputStream)
            );
            keysetString = outputStream.toString();
            System.out.println("Keyset in JSON format: " + keysetString);

            // Parse the JSON string
            ObjectMapper objectMapper = new ObjectMapper();
            JsonNode rootNode = objectMapper.readTree(keysetString);

            // Extract the value from the JSON structure
            String base64Value = rootNode.get("key").get(0).get("keyData").get("value").asText();
            byte[] keyBytes = Base64.getDecoder().decode(base64Value);

            System.out.println("Extracted key value: " + base64Value);

            // Create a CryptographyClient to encrypt the extracted key using Azure Key Vault
            CryptographyClient cryptoClient = new CryptographyClientBuilder()
                    .keyIdentifier(key.getId())
                    .credential(credential)
                    .buildClient();

            // Encrypt the key material using the Azure key
            EncryptResult encryptResult = cryptoClient.encrypt(
                    EncryptionAlgorithm.RSA_OAEP_256,
                    keyBytes);

            // Get the encrypted key material
            byte[] encryptedKeyBytes = encryptResult.getCipherText();
            String encryptedKeyBase64 = Base64.getEncoder().encodeToString(encryptedKeyBytes);

            System.out.println("Encrypted key (Base64): " + encryptedKeyBase64);


            // You can now use the aead primitive for encryption/decryption
            // For example:
            // byte[] ciphertext = aead.encrypt(plaintext, associatedData);
            // byte[] decrypted = aead.decrypt(ciphertext, associatedData);

        } catch (Exception e) {
            System.err.println("Error: " + e.getMessage());
            e.printStackTrace();
        }
    }
}