from suno_utils.utils.text import read_jsonl
import soundfile as sf
import pyloudnorm as pyln
import csv
import os
from joblib import Parallel, delayed
from tqdm import tqdm


"""
This script backfills the loudness metadata from a metas file on the local disk. 
We will iterate over the metas file, find all the local paths, load the audio files, calculate the loudness, also look for bad audio. 
Then we will create a new metadata file that is a dict lookup from id to loudness metadata.
This can then be merged in with existing metadata files or loaded on the side. 
We will use pyloudnorm to calculate the loudness, and we will use soundfile to load the audio files. 
We need to support both opus and mp3 files.
"""


def read_audio_and_calculate_loudness(meta_entry: dict) -> dict:
    local_filepath = meta_entry["local_filepath"]
    file_id = meta_entry["id"]  # get id directly from the meta entry

    result = {
        "id": file_id,
        "local_filepath": local_filepath,
        "lufs_db": None,
        "failed": False,
        "duration_s": None,
        "sample_rate": None,
    }
    try:
        audio, sr = sf.read(local_filepath)
        result["duration_s"] = float(audio.shape[0] / sr)
        result["sample_rate"] = sr
    except Exception as e:
        print(f"Error reading audio file {local_filepath}: {e}")
        result["failed"] = True
        return result

    # create a meter and measure the loudness
    try:
        meter = pyln.Meter(sr)
        # ensure the audio is stereo
        lufs_db = meter.integrated_loudness(audio)
        result["lufs_db"] = lufs_db
    except Exception as e:
        print(f"Error calculating loudness for audio file {local_filepath}: {e}")
        result["failed"] = True
        return result

    return result


def get_existing_ids(results_filepath: str) -> set:
    """Read existing CSV file and return set of already processed IDs"""
    existing_ids = set()
    if os.path.exists(results_filepath):
        try:
            with open(results_filepath, "r", newline="") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    existing_ids.add(row["id"])
            print(
                f"Found {len(existing_ids)} already processed files in {results_filepath}"
            )
        except Exception as e:
            print(
                f"Warning: Could not read existing results file {results_filepath}: {e}"
            )
            print("Starting fresh...")
    else:
        print(
            f"No existing results file found at {results_filepath}, starting fresh..."
        )
    return existing_ids


def main(metas_filepath: str, results_filepath: str, batch_size: int = 10000):
    metas = read_jsonl(metas_filepath)
    print(f"Total metas: {len(metas)}")

    # Get already processed IDs for resumability
    existing_ids = get_existing_ids(results_filepath)

    # Filter out already processed metas
    metas_to_process = [meta for meta in metas if meta["id"] not in existing_ids]
    print(f"Already processed: {len(existing_ids)}")
    print(f"Remaining to process: {len(metas_to_process)}")

    if len(metas_to_process) == 0:
        print("All files have already been processed!")
        return

    failed_count = 0
    csv_fieldnames = [
        "id",
        "local_filepath",
        "lufs_db",
        "failed",
        "duration_s",
        "sample_rate",
    ]

    # Write CSV header only if file doesn't exist
    if not os.path.exists(results_filepath):
        with open(results_filepath, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=csv_fieldnames)
            writer.writeheader()

    # Process in batches to manage memory
    for i in range(0, len(metas_to_process), batch_size):
        batch = metas_to_process[i : i + batch_size]
        batch_num = i // batch_size + 1
        total_batches = (len(metas_to_process) - 1) // batch_size + 1
        print(f"Processing batch {batch_num}/{total_batches} ({len(batch)} files)")

        # Use tqdm with the batch to provide progress for each file
        batch_results = Parallel(n_jobs=-1, backend="loky")(
            delayed(read_audio_and_calculate_loudness)(meta)
            for meta in tqdm(batch, desc=f"Batch {batch_num}", unit="file")
        )

        # Append results to CSV
        with open(results_filepath, "a", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=csv_fieldnames)
            for result in batch_results:
                writer.writerow(result)
                if result["failed"]:
                    failed_count += 1

        print(
            f"Batch {batch_num} complete. Failed: {sum(1 for r in batch_results if r['failed'])}/{len(batch)}"
        )

    print(f"Total processed in this run: {len(metas_to_process)}")
    print(f"Total failed in this run: {failed_count}/{len(metas_to_process)}")
    print(f"Total files in dataset: {len(metas)}")
    print(
        f"Total completed (including previous runs): {len(existing_ids) + len(metas_to_process) - failed_count}"
    )
    if batch_results:
        print("Sample result:", batch_results[0])


if __name__ == "__main__":
    metas_filepath = "/app2/suno/data/diffusion/v1/metas_diff_v0_tr.jsonl"
    results_filepath = "/app2/suno/data/diffusion/v1/metas_diff_v0_tr_loudness.csv"
    main(metas_filepath, results_filepath)