import os
import glob
import uuid
import torch
import json
import whisper
import torchaudio
import pyloudnorm as pyln
from tqdm import tqdm

SAMPLE_RATE = 48_000
NUM_FRAMES = int(10 * SAMPLE_RATE)
TARGET_LUFS_DB = -16.0
DRY_RUN = False
TRANSCRIBE = True

if TRANSCRIBE:
    # load whisper model
    whisper_model = whisper.load_model("base", device="cuda")

if __name__ == "__main__":
    root_dir = "/app/suno/christian/data/tiktok_covers"
    output_dir = "/app/suno/christian/data/tiktok_covers_lyrics_48khz"
    os.makedirs(output_dir, exist_ok=True)

    filepaths = sorted(glob.glob(os.path.join(root_dir, "*.mp3")))
    train_filepaths = filepaths[: int(0.9 * len(filepaths))]
    val_filepaths = filepaths[int(0.9 * len(filepaths)) :]

    meter = pyln.Meter(SAMPLE_RATE)  # create BS.1770 meter

    for subset_name, filepaths in [("val", val_filepaths), ("train", train_filepaths)]:
        subset_ouput_dir = os.path.join(output_dir, subset_name)
        metas_filepath = os.path.join(output_dir, f"{subset_name}_metas.json")
        metas = {}

        os.makedirs(subset_ouput_dir, exist_ok=True)
        for filepath in tqdm(filepaths):
            filename = os.path.basename(filepath).replace(".mp3", "")
            # use sanitized uuid
            uid = uuid.uuid4()
            print(filename)
            audio, sr = torchaudio.load(filepath)
            if sr != SAMPLE_RATE:
                audio = torchaudio.functional.resample(audio, sr, SAMPLE_RATE)

            # chunk the audio file into 10s segments
            chunk_size = 10 * 48000
            num_chunks = audio.shape[-1] // NUM_FRAMES
            for i in range(num_chunks):
                chunk = audio[..., i * NUM_FRAMES : (i + 1) * NUM_FRAMES]

                # check chunk for silence and skip if silent
                chunk_power = torch.mean(chunk**2)
                if chunk_power < 1e-4:
                    continue

                # loudness normalize this chunk
                loudness = meter.integrated_loudness(chunk.numpy().T)
                # print(loudness)
                gain_db = TARGET_LUFS_DB - loudness
                gain_lin = 10 ** (gain_db / 20)
                chunk = chunk * gain_lin
                # print(loudness, gain_db)

                if not DRY_RUN:
                    example_id = f"{uid}_{i}"
                    out_filepath = os.path.join(subset_ouput_dir, f"{example_id}.wav")
                    torchaudio.save(
                        out_filepath,
                        chunk,
                        SAMPLE_RATE,
                    )

                if TRANSCRIBE:
                    chunk_mono_16khz = torchaudio.functional.resample(
                        chunk.mean(dim=0), SAMPLE_RATE, 16_000
                    )
                    result = whisper_model.transcribe(chunk_mono_16khz)
                    lyrics = result["text"]
                    # if length of lyrics is less than 10 characters, skip
                    if len(lyrics) < 10:
                        os.remove(out_filepath)
                        continue

                    # save lyrics to file
                    metas[example_id] = {
                        "audio_file": out_filepath,
                        "title": filename,
                        "lyrics": lyrics,
                    }
                    print(metas[example_id])
        # save metas
        with open(metas_filepath, "w") as f:
            json.dump(metas, f)