import shutil
import argparse
from pathlib import Path


def organize_audio_files(source_folder, output_folder):
    """
    Organizes audio files into subfolders by song name and renames files to just instrument names.

    Expected format: Artist_-_SongName___instrument.wav
    Creates: output_folder/Artist - SongName/instrument.wav
    """

    source_path = Path(source_folder)
    output_path = Path(output_folder)

    if not source_path.exists():
        print(f"Error: Source folder '{source_folder}' does not exist.")
        return

    # Create the output folder if it doesn't exist
    output_path.mkdir(parents=True, exist_ok=True)

    # Get all .wav files in the source directory
    wav_files = list(source_path.glob("*.wav"))

    if not wav_files:
        print("No .wav files found in the source directory.")
        return

    print(f"Found {len(wav_files)} .wav files to organize...")

    for file_path in wav_files:
        filename = file_path.name

        try:
            # Parse the filename: Artist_-_SongName___instrument.wav
            if "___" not in filename:
                print(f"Skipping '{filename}' - doesn't match expected format")
                continue

            # Split by the triple underscore to separate song info from instrument
            song_part, instrument_part = filename.rsplit("___", 1)

            # Extract artist and song name
            if "_-_" in song_part:
                artist, song_name = song_part.split("_-_", 1)
                # Create folder name as "Artist - SongName"
                folder_name = f"{artist} - {song_name}"
            else:
                # Fallback if format is different
                folder_name = song_part

            # Create subfolder for this song inside the output folder
            song_folder = output_path / folder_name
            song_folder.mkdir(exist_ok=True)

            # Create new filename with just the instrument name
            new_filename = instrument_part  # This already includes .wav extension
            destination = song_folder / new_filename

            if destination.exists():
                print(f"File already exists: {destination}")
                continue

            shutil.move(str(file_path), str(destination))
            print(
                f"Moved: {filename} -> {output_path.name}/{folder_name}/{new_filename}"
            )

        except Exception as e:
            print(f"Error processing '{filename}': {str(e)}")
            continue

    print("\nOrganization complete!")

    # Display the resulting structure
    print(f"\nResulting folder structure in '{output_path}':")
    for song_folder in sorted(output_path.iterdir()):
        if song_folder.is_dir():
            print(f"  {song_folder.name}/")
            for file in sorted(song_folder.glob("*.wav")):
                print(f"    {file.name}")


def main():
    parser = argparse.ArgumentParser(
        description="Organize audio files into subfolders by song name",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Example usage:
  python organize_audio.py /path/to/source/folder /path/to/output/folder
  python organize_audio.py ./audio_files ./organized_audio
        """,
    )

    parser.add_argument(
        "source_folder", help="Path to the folder containing audio files to organize"
    )
    parser.add_argument(
        "output_folder", help="Path to the folder where organized files will be placed"
    )

    args = parser.parse_args()

    print(f"Source folder: {args.source_folder}")
    print(f"Output folder: {args.output_folder}")
    print()

    organize_audio_files(args.source_folder, args.output_folder)


if __name__ == "__main__":
    main()