#!/usr/bin/env /usr/openv/pdde/pdshared/bin/python3
# $Copyright: Copyright (c) 2022 Veritas Technologies LLC. All rights reserved $
from __future__ import absolute_import
from system_utils import ssl
ssl.enable_fips_140_2()

import os, sys
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import wait

def get_endpoint(service_host):
    if service_host.startswith('blob.'):
        return "https://{}.{}".format(os.getenv("AZURE_STORAGE_ACCOUNT"), service_host)
    else:
        return "https://{}.blob.{}".format(os.getenv("AZURE_STORAGE_ACCOUNT"), service_host)

def vhd_page_upload(service_host, container_name, blob_name, count):
    # Instantiate a new BlobServiceClient
    blob_service_client = BlobServiceClient(account_url=get_endpoint(service_host),
            credential=os.environ.get('AZURE_STORAGE_KEY'))

    # Instantiate a new ContainerClient
    container_client = blob_service_client.get_container_client(container_name)

#    try:
#        # Create new container in the service
#        container_client.create_container()
#    except ResourceExistsError as exc:
#        pass

    # Instantiate a new BlobClient
    blob_client = container_client.get_blob_client(blob_name)

    count = int(count)
    # Create the Page Blob
    blob_client.create_page_blob(count)
    stream = sys.stdin.buffer
    off = 0
    blksize = 4*1024*1024
    max_workers = 10
    executor = ThreadPoolExecutor(max_workers=max_workers)
    workers = set()
    while off < count:
        buf = stream.read(blksize)
        if buf == b'\0'*len(buf):
            off += len(buf)
            continue
        if len(workers) >= max_workers:
            # Wait for any workers to finish before submitting next one
            _, workers = wait(workers, return_when='FIRST_COMPLETED')
        # Upload content to the Page Blob
        workers.add(executor.submit(blob_client.upload_page, buf, offset=off, length=len(buf)))
        off += len(buf)
    executor.shutdown()

def blob_delete(service_host, container_name, blob_name):
    # Instantiate a new BlobServiceClient
    blob_service_client = BlobServiceClient(account_url=get_endpoint(service_host),
            credential=os.environ.get('AZURE_STORAGE_KEY'))

    # Instantiate a new ContainerClient
    container_client = blob_service_client.get_container_client(container_name)

    # Instantiate a new BlobClient
    blob_client = container_client.get_blob_client(blob_name)

    # Delete the Page Blob
    blob_client.delete_blob()

if sys.argv[1] == '-upload':
    vhd_page_upload(*sys.argv[2:])
elif sys.argv[1] == '-delete':
    blob_delete(*sys.argv[2:])
sys.exit(0)