import io
import pickle
import re
import dataiku

def get_managed_folder_info(project, managed_folder_name):
    folder_index = 0
    managed_folders = project.list_managed_folders()
    for index, folder in enumerate(managed_folders):
        if folder['name'] == managed_folder_name:
            folder_index = index
    folder_info = managed_folders[folder_index]
    return folder_info


def get_managed_folder_id_with_folder_name(project, managed_folder_name):
    folder_info = get_managed_folder_info(project, managed_folder_name)
    return folder_info["id"]


def remove_pickle_extension(pickle_name):
    import re
    return re.sub("\.p$", "", pickle_name)


def write_pickle_in_dss_folder(data, pickle_name, folder_id):
    dss_folder = dataiku.Folder(folder_id)
    pickle_bytes = io.BytesIO()
    pickle.dump(data, pickle_bytes)
    pickle_name = remove_pickle_extension(pickle_name)
    with dss_folder.get_writer(f"{pickle_name}.p") as w:
        w.write(pickle_bytes.getvalue())
        pass
    pass


def read_pickle_from_dss_folder(pickle_name, folder_id):
    import pickle
    
    dss_folder = dataiku.Folder(folder_id)
    pickle_name = remove_pickle_extension(pickle_name)
    with dss_folder.get_download_stream(f"{pickle_name}.p") as f:
        data = f.read()
        pickle_data = pickle.loads(data)
    f.close()
    return pickle_data