from .Writer import Writer
from azure.storage.blob import BlockBlobService
import json
import pandas as pd
import numpy as np


class ADLSWriter(Writer):
    """
        Write dataframe/json to specified blob storage location
    """
    def __init__(self, account_name, account_key, container_name, storage_name, dataflow_name):
        self.account_name = account_name
        self.account_key = account_key
        self.container_name = container_name
        self.storage_name = storage_name
        self.dataflow_name = dataflow_name
    
    def write_df(self, blob_location, dataframe, number_of_partition=5):
        """
            Write dataframe to specified blob storage location
        """
        block_blob_service = BlockBlobService(account_name=self.account_name, account_key=self.account_key)
        dfs = np.array_split(dataframe, number_of_partition)
        result = list()

        entity_name = blob_location.split('/')[0]
        blob_location = blob_location + "/" + entity_name
        for i in range(len(dfs)):
            dataframe = dfs[i].to_csv(index=False, header=False)
            filename = blob_location + str(i) + ".csv"
            block_blob_service.create_blob_from_text(self.container_name + "/" + self.dataflow_name, 
                                                     filename, dataframe)
            blob_url = 'https://' + self.storage_name + '.dfs.core.windows.net/' + self.container_name + '/' + filename
            result.append((filename, blob_url))
        return result

    def write_json(self, blob_location, json_dict):
        """
        write json to specified blob storage location
        """
        json_dict = json.dumps(json_dict)
        block_blob_service = BlockBlobService(
            account_name=self.account_name, account_key=self.account_key)
        block_blob_service.create_blob_from_text(self.container_name,self.dataflow_name+"/"+ blob_location, json_dict)
        blob_url = 'https://'+self.storage_name+'.dfs.core.windows.net/'+self.container_name+'/'+blob_location
        return blob_url
