# %%
from suno_utils.audio import Audio
import requests
import os

# %%
params = {
    "gen_id": "f62ab88d-0d12-4d31-bc7a-89b36730b6c5",
    "min_loop_length_bars": 2,
    "max_loop_length_bars": 4,
    #"override_target_tempo": 144,
    #"confidence_cutoff": 0.8,
    "output_stems": True,
}

response = requests.get(
    "https://suno-ai--loop-extraction-dev-extract-loop-points.modal.run",
    params=params
)

# %%
response_json = response.json()
gen_id = response_json['gen_id']
loops = response_json['loops']
print(gen_id)
print(f"Found {len(loops)} loops")

# %%
audio_paths = []
response_json = response.json()
for loop, loop_data in response_json["loops"].items():
    bars = loop_data["bar_length"]
    tempo = loop_data["tempo"]
    print(f"{bars} bars, {tempo} bpm")
    if "s3_path" in loop_data:
        parent_loop = loop_data["s3_path"]
        print(f"Parent: {parent_loop}")
        audio_paths.append(parent_loop)
    if "stems" in loop_data:
        stems = loop_data["stems"]
        stem_loudness = loop_data["loudness_metrics"]
        for stem, loudness in zip(stems, stem_loudness):
            if loudness and loudness["peak_db"] > -10:
                print(stem, loudness["peak_db"])
                audio_paths.append(stem)

print(audio_paths)

# %%
audios = [Audio.from_s3(ap) for ap in audio_paths]

for audio in audios:
    audio_loop = Audio.concatenate([audio, audio, audio])
    audio_loop.play()

# %%
output_dir = f"/home/sara/sara/{gen_id}"
os.makedirs(output_dir, exist_ok=True)

for idx, audio in enumerate(audios):
    local_fp = os.path.join(output_dir, f"{idx}.wav")
    audio.write_wav(local_fp)