import os
import glob
import json
import labelbox as lb

client = lb.Client(api_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOiJjbHppbmJnY2wwMDYyMDd5bWg2enhiMTd6Iiwib3JnYW5pemF0aW9uSWQiOiJjbHppbmJnY2QwMDYxMDd5bWM4cjM5cHdzIiwiYXBpS2V5SWQiOiJjbHp3cmU3ZnQwYjFzMDd6aWRrNTFnb21mIiwic2VjcmV0IjoiMGU5M2MwNmU4ZWI1Y2Y3NTlmNTk5YTk5MzIwOTU5MzQiLCJpYXQiOjE3MjM4MTU3NTcsImV4cCI6MjM1NDk2Nzc1N30.Z6gZwlzQ85KrqGOydqCdo1RVmUhOEntpev2HhNso4PU")
dataset = client.create_dataset(name="covers-test")


# find all examples 
base_dirname = "covers-20240816"
local_dir = f"/home/christian/code/christian/notebooks/outputs/{base_dirname}"
s3_bucket = f"s3://suno-annotation-public/{base_dirname}"
s3_bucket_url = f"https://suno-annotation-public.s3.amazonaws.com/{base_dirname}"

# find all json files 
metas = glob.glob(os.path.join(local_dir, "*.json"))
print(f"Found {len(metas)} examples...")

metas = metas[:10] # just the first 10 for now

assets = [] # list of assets to add to the dataset

for meta in metas:
    with open(meta, "r") as f:
        data = json.load(f)

    meta_id = data["uid"]

    # upload the audio files to s3 bucket
    os.system(f"aws s3 cp {local_dir}/{meta_id}-source.mp3 {s3_bucket}/")
    os.system(f"aws s3 cp {local_dir}/{meta_id}-cover-0.mp3 {s3_bucket}/")
    os.system(f"aws s3 cp {local_dir}/{meta_id}-cover-1.mp3 {s3_bucket}/")

    # build paths to audio files with this uid
    s3_src_path = os.path.join(s3_bucket_url, f"{meta_id}-source.mp3")
    s3_gen_a_path = os.path.join(s3_bucket_url, f"{meta_id}-cover-0.mp3")
    s3_gen_b_path = os.path.join(s3_bucket_url, f"{meta_id}-cover-1.mp3")

    target_tags = data["target_tags"]
    print("target tags:", target_tags)
    lyrics = data["lyrics"]

    assets.append({
    "row_data": f"cover-example-{meta_id}",
    "global_key": f"cover-example-{meta_id}",
    "media_type": "TEXT",
    #"metadata_fields": [{"schema_id": "cko8s9r5v0001h2dk9elqdidh", "value": "my tag"}],
    "attachments":  [{"type": "RAW_TEXT", "value": ", ".join(target_tags)},
                     {"type": "RAW_TEXT", "value": lyrics},
                    {"type": "AUDIO", "value": s3_src_path},
                    {"type": "AUDIO", "value":  s3_gen_a_path},
                    {"type": "AUDIO", "value": s3_gen_b_path}],
    })

# Bulk add data rows to the dataset
print(assets[1])
#task = dataset.create_data_rows(assets)
#task.wait_till_done()
#print(task.errors)