# fixes poison directory bug
import os
import argparse

if __name__ == "__main__":
    del_str = f'[a-z,0-9]'
    is_gs = lambda x: x.startswith("gs://")


    def bucket_name(path):
        if is_gs:
            return path.replace("gs://", "").split("/")[0]
        else:
            raise ValueError(f"{path} is not a cloud bucket path")


    parser = argparse.ArgumentParser()
    parser.add_argument("directory", type=str, help="Name of directory to fix")
    parser.add_argument("-y", "--yes", action="store_true", help="pass this flag to ignore warning")
    args = parser.parse_args()

    if not args.yes:
        assert is_gs(args.directory)
        resp = input(f"WARNING: this script will delete any files in {args.directory} that don't begin "
                     f"with any of the following characters: {del_str}. Continue?")
    else:
        resp = "y"
    if resp.lower().strip() in ["y", "yes"]:
        if not args.directory.endswith("/"):
            args.directory += "/"
        bucketname = bucket_name(args.directory)
        cmd = f"gsutil cp {args.directory}{del_str}* gs://{bucketname}/tmp"
        os.system(cmd)
        cmd = f"gsutil rm -rf {args.directory}"
        os.system(cmd)
        cmd = f"gsutil mv gs://{bucketname}/tmp/* {args.directory}"
        os.system(cmd)