import psutil
import random
import requests
import subprocess
import time

if __name__ == "__main__":
    # run seaweed server
    seaweed_process = subprocess.Popen(["weed", "server", "-volume.max=50000", "-dir=./data", "-ip=127.0.0.1"])
    pid = seaweed_process.pid

    # make 1MB blob to insert
    blob  = "".join(str(random.randint(0, 255)) for i in range(400000))
    with open("test.blob", "w") as f:
        f.write(blob)

    # wait 30s for the seaweed server to start up
    time.sleep(30)

    # make the log file
    log_filename = "memory_usage_collections.log"
    with open(log_filename, "w") as f:
        f.write("timestamp, memory usage\n")

    for i in range(int(1e4)):
        # assign file key
        collection_name = f"collection{i}"
        r = requests.get(f"http://127.0.0.1:9333/dir/assign?collection={collection_name}")
        volume_id = r.json()["fid"]

        # insert blob into volume
        # curl -F file=test.blob http://127.0.0.1:8080/volume_id
        files = {"file": (None, "test.blob"),}
        r = requests.post(f"http://127.0.0.1:8080/{volume_id}", files=files)
        
        # only log memory usage each minute or so
        if i % 60 == 0:
            
            # get memory usage of seaweed process
            process = psutil.Process(pid)
            mem = process.memory_info()
            mem_time = time.time()
            
            # write to file
            with open(log_filename, "a") as f:
                f.write(f"{mem_time},{mem}\n")
        
        time.sleep(1)