import os
import itertools
import random
from pathlib import Path

BASE_NAME = "vox"


def get_song_names(base_path):
    """Get all song names from the ground_truth folder."""
    ground_truth_path = Path(base_path) / f"original_format_trimmed_{BASE_NAME}"
    if not ground_truth_path.exists():
        raise FileNotFoundError(f"Ground truth folder not found: {ground_truth_path}")

    song_names = []
    for item in ground_truth_path.iterdir():
        if item.is_dir():
            song_names.append(item.name)

    return sorted(song_names)


def verify_vocals_exist(base_path, model_folders, song_names):
    """Verify that vocals.wav exists for all songs in all model folders."""
    missing_files = []

    for model in model_folders:
        model_path = Path(base_path) / model
        for song in song_names:
            vocals_path = model_path / song / f"{BASE_NAME}.wav"
            if not vocals_path.exists():
                missing_files.append(str(vocals_path))

    if missing_files:
        print("Warning: Missing vocals.wav files:")
        for file in missing_files:
            print(f"  - {file}")
        print()

    return len(missing_files) == 0


def create_vocal_pairs(
    base_path, output_file=f"{BASE_NAME}_pairs.txt", random_seed=None
):
    """
    Create all pairs of vocals.wav between different models with random ordering.

    Args:
        base_path: Path to the parent folder containing model folders
        output_file: Name of the output file to save pairs
        random_seed: Seed for random number generator (for reproducible results)
    """

    if random_seed is not None:
        random.seed(random_seed)

    base_path = Path(base_path)

    # Define model folders
    model_folders = [
        f"original_format_trimmed_{BASE_NAME}",
        f"lalal_format_trimmed_{BASE_NAME}",
        f"suno_format_trimmed_{BASE_NAME}",
    ]

    # Verify all model folders exist
    for model in model_folders:
        model_path = base_path / model
        if not model_path.exists():
            raise FileNotFoundError(f"Model folder not found: {model_path}")

    print(f"Base path: {base_path}")
    print(f"Model folders: {model_folders}")

    # Get song names from ground_truth folder
    song_names = get_song_names(base_path)
    print(
        f"Found {len(song_names)} songs: {song_names[:5]}{'...' if len(song_names) > 5 else ''}"
    )

    # Verify all vocals.wav files exist
    all_files_exist = verify_vocals_exist(base_path, model_folders, song_names)
    if not all_files_exist:
        response = input("Some vocals.wav files are missing. Continue anyway? (y/n): ")
        if response.lower() != "y":
            return

    # Generate all pairs between different models
    pairs = []

    for song in song_names:
        # Create all combinations of models (pairs)
        model_combinations = list(itertools.combinations(model_folders, 2))

        for model1, model2 in model_combinations:
            path1 = base_path / model1 / song / f"{BASE_NAME}.wav"
            path2 = base_path / model2 / song / f"{BASE_NAME}.wav"

            # Skip if either file doesn't exist
            if not (path1.exists() and path2.exists()):
                continue

            # Randomly decide the order of the pair
            if random.choice([True, False]):
                pairs.append((str(path1), str(path2)))
            else:
                pairs.append((str(path2), str(path1)))

    # Shuffle all pairs
    random.shuffle(pairs)

    # Write pairs to file
    output_path = Path(output_file)
    with open(output_path, "w") as f:
        # Write CSV header
        f.write(
            "song_name,source_a,source_b,source_a_fp,source_b_fp,mixture_path,instrument\n"
        )

        for file1, file2 in pairs:
            # Extract song name and source model from paths
            path1 = Path(file1)
            path2 = Path(file2)

            song_name = path1.parent.name  # Should be same for both
            source_a = path1.parent.parent.name  # Model folder name
            source_b = path2.parent.parent.name  # Model folder name

            # Create path to ground truth mixture.wav
            mixture_path = (
                base_path
                / f"original_format_trimmed_{BASE_NAME}"
                / song_name
                / f"mixture_{BASE_NAME}_trimmed.wav"
            )

            instrumental = "instrumental" if BASE_NAME == "vox" else BASE_NAME

            f.write(
                f"{song_name},{source_a},{source_b},{file1.replace('vox.wav', 'instrumental.wav')},{file2.replace('vox.wav', 'instrumental.wav')},{mixture_path},{instrumental}\n"
            )

    print(f"\nGenerated {len(pairs)} pairs")
    print(f"Results saved to: {output_path.absolute()}")

    return pairs


def main():
    # Example usage
    base_path = (
        "/app2/suno/data/sara/musdb/audio_comparison2/"  # Update this path as needed
    )

    # You can set a seed for reproducible results
    pairs = create_vocal_pairs(
        base_path=base_path,
        output_file=os.path.join(base_path, f"{BASE_NAME}_pairs_trimmed.txt"),
        random_seed=42,  # Remove or set to None for truly random results
    )

    if pairs:
        print("\nFirst 3 pairs as example:")
        for i, (file1, file2) in enumerate(pairs[:3], 1):
            print(
                f"  Pair {i}: {Path(file1).parent.parent.name} vs {Path(file2).parent.parent.name}"
            )
            print(f"    Song: {Path(file1).parent.name}")


if __name__ == "__main__":
    main()