import os
import zipfile
import json
import shutil
import math

print("Looking for Note Block Studio datapack export ZIP files...")
scanned_zip_files = [file for file in os.listdir() if file.endswith(".zip")]
if len(scanned_zip_files) == 0:
    input("No zip files found!")
    quit()

# Ensure temp directory is clean
if os.path.exists("temp_extract"):
    shutil.rmtree("temp_extract")
os.makedirs("temp_extract", exist_ok=True)

# Dump everything into the temp folder
for filename in scanned_zip_files:
    print(f"Loading {filename}")
    with zipfile.ZipFile(filename, 'r') as zip:
        for item in zip.namelist():
            if item.startswith("data/") and not item.startswith("data/minecraft") and not item.endswith("/"):
                zip.extract(item, "temp_extract")

print(f"Preparing to export {len(scanned_zip_files)} songs")

for namespace in os.listdir(os.path.join("temp_extract", "data")):
    song = {}
    song["name"] = namespace
    song["weight"] = 10
    song["credits"] = ""
    song["notes"] = []

    with open(os.path.join("temp_extract", "data", namespace, "functions", "load.mcfunction"), "r", encoding="utf-8") as tempo_file:
        for line in tempo_file:
            if line.startswith("scoreboard players set speed nbs_"):
                # Convert tempo from scoreboard ticks to ticks per second
                tempo = int(line.split()[-1]) / 4
                # Write tempo as a ticks-per-note multiplier
                song["tempo"] = round(20 / tempo)

    note_index = -1

    for note_file in sorted(
        os.listdir(os.path.join("temp_extract", "data", namespace, "functions", "notes")),
        key=lambda name: int(os.path.splitext(name)[0])
    ):
        previous_index = note_index
        note_index = int(note_file.replace(".mcfunction", ""))
        with open(os.path.join("temp_extract", "data", namespace, "functions", "notes", note_file)) as file:
            # Convert pitch from multiplier to Supplementaries keys
            try:
                note_pitch_multiplier = float(file.readline().split()[-2])
            except ValueError:
                if os.path.exists("temp_extract"):
                    shutil.rmtree("temp_extract")
                input(f"Note {note_index} seems to have no sound. Aborting.")
                quit()
            note_pitch_key = round(1 + 12 * math.log2(2 * note_pitch_multiplier))
            # Write note pitch and duration
            if previous_index > -1: song["notes"].append("-"+str(note_index - previous_index))
            song["notes"].append(note_pitch_key)
    # Repeat last note timing and hope for the best
    song["notes"].append("-"+str(note_index - previous_index))

    with open(namespace+".json", "w", encoding="utf-8") as file:
        json.dump(song, file, indent=2)
    print(f"Exported song \"{namespace}.json\"")

shutil.rmtree("temp_extract")
input(f"Done! Created {len(scanned_zip_files)} song files")