import static org.assertj.core.api.Assertions.assertThat;

import com.activeviam.cloud.test.internal.util.impl.CloudTestUtil;
import com.azure.core.cryptography.AsyncKeyEncryptionKey;
import com.azure.core.http.rest.PagedIterable;
import com.azure.security.keyvault.keys.cryptography.KeyEncryptionKeyClientBuilder;
import com.azure.security.keyvault.keys.models.JsonWebKey;
import com.azure.security.keyvault.keys.models.KeyOperation;
import com.azure.storage.blob.BlobContainerClient;
import com.azure.storage.blob.BlobServiceClient;
import com.azure.storage.blob.BlobServiceClientBuilder;
import com.azure.storage.blob.models.BlobItem;
import com.azure.storage.blob.models.ListBlobsOptions;
import com.azure.storage.blob.specialized.cryptography.EncryptedBlobClient;
import com.azure.storage.blob.specialized.cryptography.EncryptedBlobClientBuilder;
import com.azure.storage.blob.specialized.cryptography.EncryptionVersion;
import java.io.ByteArrayInputStream;
import java.nio.charset.Charset;
import java.security.KeyPair;
import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

public class TestAzureUpdateWithEncryptedDirectory {

  private static final String connectionString = getConnectionString();

  private static BlobContainerClient CONTAINER;
  private static final String CONTAINER_NAME = "TestAzureUpdateWithEncryptedDirectory".toLowerCase();
  private static AsyncKeyEncryptionKey keyEncryptionKey;
  private static String keyEncryptionAlgorithm;

  @BeforeAll
  public static void setUp() throws Exception {
    BlobServiceClient client = new BlobServiceClientBuilder().connectionString(connectionString).buildClient();
    initializeEncryptionCapabilities(
        TestAzureUpdateWithEncryptedDirectory.class.getName(),
        CloudTestUtil.getEncryptionKey(),
        CloudTestUtil.getKeyEncryptionAlgorithm());
    CONTAINER = client.getBlobContainerClient(CONTAINER_NAME);
    CONTAINER.deleteIfExists(); // useful when rerunning the test to clean up the data from previous tests
    Thread.sleep(Duration.ofSeconds(5)); // wait for deletion of the container
    CONTAINER.create();
  }

  @Test
  void testListEntitiesOnAzure() {
    final String e1 = "entity1";
    final String e2 = "dir1/entity2";
    PagedIterable<BlobItem> blobs;

    blobs = CONTAINER.listBlobs(new ListBlobsOptions().setPrefix(""), null);
    assertThat(blobs.stream().map(BlobItem::getName)).isEmpty();

    uploadData(e1, "1");
    blobs = CONTAINER.listBlobs(new ListBlobsOptions().setPrefix(""), null);
    assertThat(blobs.stream().map(BlobItem::getName)).containsExactlyInAnyOrder(e1);

    uploadData(e2, "2");
    blobs = CONTAINER.listBlobs(new ListBlobsOptions().setPrefix(""), null);
    assertThat(blobs.stream().map(BlobItem::getName)).containsExactlyInAnyOrder(e1, e2);
  }


  private void uploadData(final String name, final String content) {
    try {

      final byte[] buffer = content.getBytes(Charset.defaultCharset());
      final EncryptedBlobClient encryptionClient =
          new EncryptedBlobClientBuilder(EncryptionVersion.V2)
              .blobClient(CONTAINER.getBlobClient(name))
              .key(keyEncryptionKey, keyEncryptionAlgorithm)
              .buildEncryptedBlobClient();

      encryptionClient.upload(new ByteArrayInputStream(buffer), buffer.length);

    } catch (final Exception e) {
      throw new RuntimeException(name, e);
    }
  }

  private static void initializeEncryptionCapabilities(
      final String keyId, final KeyPair keyPair, final String encryptionAlgorithm) {
    keyEncryptionKey = buildKeyEncryptionKeyFrom(keyId, keyPair, List.of(KeyOperation.UNWRAP_KEY, KeyOperation.WRAP_KEY));
    keyEncryptionAlgorithm = encryptionAlgorithm;
  }

  private static AsyncKeyEncryptionKey buildKeyEncryptionKeyFrom(
      final String keyId, final KeyPair keyPair, final List<KeyOperation> authorizedOperations) {
    final JsonWebKey jsonKey = JsonWebKey.fromRsa(keyPair, authorizedOperations).setId(keyId);
    return new KeyEncryptionKeyClientBuilder().buildAsyncKeyEncryptionKey(jsonKey).block();
  }
}