import shutil
from pathlib import Path
from collections import defaultdict


def extract_song_info(filename):
    """Extract song name and type from LALAL filename."""
    name = filename.replace(".wav", "")

    if name.endswith("_no_vocals_split_by_lalalai"):
        return name.replace("_no_vocals_split_by_lalalai", ""), "instrumental"
    elif name.endswith("_vocals_split_by_lalalai"):
        return name.replace("_vocals_split_by_lalalai", ""), "vocals"
    else:
        print(f"Warning: Unexpected format: {filename}")
        return name, "unknown"


def reformat_lalal_folder(lalal_path, dry_run=True):
    """Reformat LALAL folder to match expected structure."""
    lalal_path = Path(lalal_path)

    if not lalal_path.exists():
        raise FileNotFoundError(f"Folder not found: {lalal_path}")

    wav_files = [f for f in lalal_path.glob("*.wav") if not f.name.startswith("._")]
    if not wav_files:
        print(f"No .wav files found in {lalal_path}")
        return []

    # Group files by song
    songs = defaultdict(dict)
    for wav_file in wav_files:
        song_name, file_type = extract_song_info(wav_file.name)
        if file_type != "unknown":
            songs[song_name][file_type] = wav_file

    # Create operations list
    operations = []
    for song_name, files in songs.items():
        target_dir = lalal_path / song_name
        for file_type, source_file in files.items():
            target_filename = f"{file_type}.wav"
            operations.append(
                {
                    "source": source_file,
                    "target": target_dir / target_filename,
                    "song": song_name,
                    "type": file_type,
                }
            )

    # Show summary
    print(f"Found {len(songs)} songs with {len(operations)} files:")
    for song, files in sorted(songs.items()):
        types = ", ".join(files.keys())
        print(f"  {song}: {types}")

    if dry_run:
        print("\nDRY RUN - Set dry_run=False to execute")
        return operations

    # Execute operations
    if input(f"\nMove {len(operations)} files? (y/n): ").lower() != "y":
        return operations

    success = 0
    for op in operations:
        try:
            op["target"].parent.mkdir(exist_ok=True)
            shutil.move(str(op["source"]), str(op["target"]))
            print(f"✓ {op['song']}/{op['type']}.wav")
            success += 1
        except Exception as e:
            print(f"✗ {op['source'].name}: {e}")

    print(f"\nCompleted: {success}/{len(operations)} files moved")
    return operations


def verify_with_ground_truth(lalal_path, ground_truth_path):
    """Compare song names with ground truth folder."""
    lalal_songs = set()
    for wav_file in Path(lalal_path).glob("*.wav"):
        if not wav_file.name.startswith("._"):  # Skip macOS metadata files
            song_name, file_type = extract_song_info(wav_file.name)
            if file_type != "unknown":
                lalal_songs.add(song_name)

    gt_songs = {d.name for d in Path(ground_truth_path).iterdir() if d.is_dir()}

    missing = gt_songs - lalal_songs
    extra = lalal_songs - gt_songs

    print(f"Ground truth: {len(gt_songs)}, LALAL: {len(lalal_songs)}")
    if missing:
        print(f"Missing in LALAL: {missing}")
    if extra:
        print(f"Extra in LALAL: {extra}")
    if not missing and not extra:
        print("✓ Perfect match!")


def main():
    base_path = "/app2/suno/data/sara/musdb/audio_comparison2/audio_comparison"
    lalal_path = f"{base_path}/lalal"
    ground_truth_path = f"{base_path}/original"

    try:
        # Verify against ground truth
        if Path(ground_truth_path).exists():
            verify_with_ground_truth(lalal_path, ground_truth_path)
            print()

        # Show what would be done
        operations = reformat_lalal_folder(lalal_path, dry_run=False)

        # Uncomment to execute:
        # reformat_lalal_folder(lalal_path, dry_run=False)

    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()