STEP = 0
def step(msg):
    global STEP
    print(f"STEP {STEP}: {msg}")
    input("Press Enter to continue...")
    STEP += 1

step("Clean up any old data")

import shutil
import os
shutil.rmtree("garage-rebalance-toast", ignore_errors=True)


step("Set up Garage config")

os.makedirs("garage-rebalance-toast/meta")
os.makedirs("garage-rebalance-toast/data-old")
os.makedirs("garage-rebalance-toast/data-new1")
os.makedirs("garage-rebalance-toast/data-new2")

SECRET = os.urandom(32).hex()
config = """
metadata_dir = "garage-rebalance-toast/meta"
data_dir = [
    { path = "garage-rebalance-toast/data-old", capacity = "1G" },
]

db_engine = "lmdb"

replication_factor = 1

rpc_bind_addr = "[::]:3901"
rpc_public_addr = "127.0.0.1:3901"
rpc_secret = "SECRET"

[s3_api]
s3_region = "garage"
api_bind_addr = "[::]:3900"
root_domain = ".s3.garage.localhost"
""".replace("SECRET", SECRET)

with open("garage-rebalance-toast/garage.toml", "w") as f:
    f.write(config)

step("Run the first Garage instance")

import subprocess
import sys

def garage_run(*args, **kwargs):
    return subprocess.Popen([os.path.expanduser("~/.cargo/bin/garage"), "-c", "garage-rebalance-toast/garage.toml", *args], env={"RUST_LOG": "debug"}, **kwargs)

server_proc = garage_run("server")

import time
# Time for the server to settle
time.sleep(5)


step("Setup the layout")

node_status = garage_run("status", stdout=subprocess.PIPE).communicate(None)[0].decode("utf-8").strip()
node_id = node_status.split("\n")[-1].split(" ")[0]
print(f"Node ID: {node_id}")

garage_run("layout", "assign", "-c", "1GB", node_id, "-z", "zone1").communicate(None)
garage_run("layout", "apply", "--version", "1").communicate(None)
time.sleep(2)

step("Create an access key")
key_output = garage_run("key", "create", stdout=subprocess.PIPE).communicate(None)[0].decode("utf-8").strip()
key_id = key_output.split("\n")[1].split(":")[-1].strip()
key_secret = key_output.split("\n")[2].split(":")[-1].strip()

print(f"Key ID: {key_id}")
print(f"Secret key: {key_secret}")

garage_run("key", "allow", "--create-bucket", key_id).communicate(None)

step("Setup Rclone config")


with open("garage-rebalance-toast/rclone.conf", "w") as f:
    f.write(f"""
[garage]
type = s3
provider = Other
env_auth = false
access_key_id = {key_id}
secret_access_key = {key_secret}
endpoint = http://127.0.0.1:3900
region = garage
""")


step("Upload some data")

def rclone_run(*args, **kwargs):
    return subprocess.Popen(["rclone", *args], env={"RCLONE_CONFIG": "garage-rebalance-toast/rclone.conf"}, **kwargs)

with open("garage-rebalance-toast/randomness.bin", "wb") as f:
    f.write(os.urandom(1024 * 1024 * 50))

rclone_run("move", "-vvP", "garage-rebalance-toast/randomness.bin", "garage:bucket/randomness").communicate(None)


step("Stop Garage and reconfigure its data locations")
server_proc.terminate()

server_proc.communicate(None)

time.sleep(3)

config = """
metadata_dir = "garage-rebalance-toast/meta"
data_dir = [
    { path = "garage-rebalance-toast/data-old", read_only = true },
    { path = "garage-rebalance-toast/data-new1", capacity = "1G" },
    { path = "garage-rebalance-toast/data-new2", capacity = "1G" },
]

db_engine = "lmdb"

replication_factor = 1

rpc_bind_addr = "[::]:3901"
rpc_public_addr = "127.0.0.1:3901"
rpc_secret = "SECRET"

[s3_api]
s3_region = "garage"
api_bind_addr = "[::]:3900"
root_domain = ".s3.garage.localhost"
""".replace("SECRET", SECRET)

with open("garage-rebalance-toast/garage.toml", "w") as f:
    f.write(config)


step("Copy a single file from the old location to the new one, then truncate it at the source")

first_part = list(filter(lambda x: x != "garage-marker", os.listdir("garage-rebalance-toast/data-old")))[0]
second_part = os.listdir(f"garage-rebalance-toast/data-old/{first_part}")[0]
file = os.listdir(f"garage-rebalance-toast/data-old/{first_part}/{second_part}")[0]
os.makedirs(f"garage-rebalance-toast/data-new1/{first_part}/{second_part}", exist_ok=True)

shutil.copy(f"garage-rebalance-toast/data-old/{first_part}/{second_part}/{file}", f"garage-rebalance-toast/data-new1/{first_part}/{second_part}/{file}")

# length approximately equal to observed length when doing this live
os.truncate(f"garage-rebalance-toast/data-old/{first_part}/{second_part}/{file}", 64*1024)


step("Run the second Garage instance")

server_proc = garage_run("server")

time.sleep(5)

step("Initiate a rebalance and wait 20 seconds")

garage_run("repair", "--yes", "rebalance")

time.sleep(20)

step("Check the data status")

old_files = 0
for root, dirs, files in os.walk("garage-rebalance-toast/data-old"):
    for file in files:
        old_files += 1


new1_files = 0
for root, dirs, files in os.walk("garage-rebalance-toast/data-new1"):
    for file in files:
        new1_files += 1

new2_files = 0
for root, dirs, files in os.walk("garage-rebalance-toast/data-new2"):
    for file in files:
        new2_files += 1

print(f"Old files: {old_files}")
print(f"New1 files: {new1_files}")
print(f"New2 files: {new2_files}")

print("We expected that old contains zero files and the new directories contain all of them.")

step("Finally, stop the server")
server_proc.terminate()
server_proc.communicate(None)