# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_remote_utils.ipynb.

# %% auto 0
__all__ = ['is_bucket', 'gsutil_bucket', 'gsutil_src', 'bucket_del', 'bucket_count', 'bucket_move', 'bucket_up', 'bucket_dl',
           'get_local_path', 'handle_input_path']

# %% ../nbs/02_remote_utils.ipynb 2
from dreamai.imports import *
from .imports import *
from .utils import *


# %% ../nbs/02_remote_utils.ipynb 4
def is_bucket(p):
    return str(p).startswith("gs://")


def gsutil_bucket(bucket):
    if not str(bucket).startswith("gs://"):
        bucket = "gs://" + str(bucket)
    if bucket[-1] != "/" and Path(bucket).suffix != ".json":
        bucket += "/"
    return bucket


def gsutil_src(folder):
    if Path(folder).suffix != "":
        return str(folder)
    folder = str(folder)
    if folder[-1] != "/":
        folder += "/"
    folder += "*"
    return folder


def bucket_del(bucket):
    msg.info(f"Deleting {bucket}.", spaced=True)
    gu = shutil.which("gsutil")
    bucket = gsutil_bucket(bucket)
    subprocess.run([gu, "-m", "rm", "-r", bucket])


def bucket_count(bucket):
    gu = shutil.which("gsutil")
    bucket = gsutil_src(bucket)
    cmd = [gu, "ls", "-l", bucket]
    res = subprocess.run(cmd, stdout=subprocess.PIPE)
    res = res.stdout.decode("utf-8").split("\n")
    return len(res) - 2


def bucket_move(folder, bucket):
    msg.info(f"Moving {folder} to {bucket}.", spaced=True)
    gu = shutil.which("gsutil")
    bucket = gsutil_bucket(bucket)
    folder = gsutil_src(folder)
    files = get_files(folder, make_str=True)
    if len(files) == 1:
        folder = files[0]
    subprocess.run([gu, "-m", "mv", folder, bucket])


def bucket_up(folder, bucket, only_new=True):
    msg.info(f"Uploading {folder} to {bucket}.", spaced=True)
    gu = shutil.which("gsutil")
    bucket = gsutil_bucket(bucket)
    folder = gsutil_src(folder)
    cmd = [gu, "-m", "cp"]
    if only_new:
        cmd.append("-n")
    files = get_files(folder, make_str=True)
    if len(files) == 1:
        folder = files[0]
    else:
        cmd.append("-r")
    subprocess.run(cmd + [folder, bucket])


def bucket_dl(bucket, folder, only_new=True):
    msg.info(f"Downloading {bucket} to {folder}.", spaced=True)
    gu = shutil.which("gsutil")
    bucket = gsutil_bucket(bucket)
    bucket = gsutil_src(bucket)
    os.makedirs(folder, exist_ok=True)
    cmd = [gu, "-m", "cp", "-r"]
    if only_new:
        cmd.append("-n")
    subprocess.run(cmd + [bucket, folder])


def get_local_path(remote_path, local_path):
    if is_bucket(remote_path):
        if Path(remote_path).suffix != "":
            return Path(local_path), remote_path
        return Path(local_path) / Path(remote_path).name, remote_path
    else:
        return Path(remote_path), remote_path


def handle_input_path(path, local_path=None, task_id=gen_random_string(16), only_new=True):
    path = str(path)
    if not is_bucket(path):
        return path, path
    if local_path is None:
        local_path = Path(f"/tmp/{task_id}")
    os.makedirs(local_path, exist_ok=True)
    local_path, bucket = get_local_path(path, local_path)
    bucket_dl(bucket, local_path, only_new=only_new)
    return local_path, bucket
