#!/usr/bin/env python3
"""
Audio Instrumental Generator
Creates instrumental.wav files by subtracting vox.wav from mixture.wav
"""

import librosa
import soundfile as sf
import numpy as np
from pathlib import Path
import argparse


def create_instrumental(song_folder):
    """
    Create instrumental.wav by subtracting vox.wav from mixture.wav

    Args:
        song_folder (str): Path to the song folder containing stems
    """
    song_path = Path(song_folder)
    mixture_path = song_path / "mixture.wav"
    vox_path = song_path / "vox.wav"
    instrumental_path = song_path / "instrumental.wav"

    # Check if required files exist
    if not mixture_path.exists():
        print(f"Warning: {mixture_path} not found, skipping {song_path.name}")
        return False

    if not vox_path.exists():
        print(f"Warning: {vox_path} not found, skipping {song_path.name}")
        return False

    # Check if instrumental already exists
    if instrumental_path.exists():
        print(f"Instrumental already exists for {song_path.name}, skipping")
        return True

    try:
        # Load audio files
        print(f"Processing {song_path.name}...")
        mixture, sr_mix = librosa.load(str(mixture_path), sr=None, mono=False)
        vox, sr_vox = librosa.load(str(vox_path), sr=None, mono=False)

        # Ensure sample rates match
        if sr_mix != sr_vox:
            print(
                f"Warning: Sample rate mismatch in {song_path.name} (mixture: {sr_mix}, vox: {sr_vox})"
            )
            # Resample vox to match mixture
            vox = librosa.resample(vox, orig_sr=sr_vox, target_sr=sr_mix)

        # Handle mono/stereo differences
        if mixture.ndim != vox.ndim:
            if mixture.ndim == 1 and vox.ndim == 2:
                mixture = np.array([mixture, mixture])  # Convert mono to stereo
            elif mixture.ndim == 2 and vox.ndim == 1:
                vox = np.array([vox, vox])  # Convert mono to stereo

        # Ensure same length by padding or truncating
        min_length = min(mixture.shape[-1], vox.shape[-1])
        if mixture.ndim == 1:
            mixture = mixture[:min_length]
            vox = vox[:min_length]
        else:
            mixture = mixture[:, :min_length]
            vox = vox[:, :min_length]

        # Create instrumental by subtraction
        instrumental = mixture - vox

        # Save the instrumental track
        sf.write(
            str(instrumental_path),
            instrumental.T if instrumental.ndim == 2 else instrumental,
            sr_mix,
        )
        print(f"Created instrumental for {song_path.name}")
        return True

    except Exception as e:
        print(f"Error processing {song_path.name}: {str(e)}")
        return False


def process_music_database(root_folder):
    """
    Process all song folders in the music database

    Args:
        root_folder (str): Root folder containing song subfolders
    """
    root_path = Path(root_folder)

    if not root_path.exists():
        print(f"Error: Root folder {root_folder} does not exist")
        return

    song_folders = [f for f in root_path.iterdir() if f.is_dir()]

    if not song_folders:
        print(f"No subfolders found in {root_folder}")
        return

    print(f"Found {len(song_folders)} song folders")
    processed = 0
    skipped = 0
    errors = 0

    for song_folder in sorted(song_folders):
        result = create_instrumental(song_folder)
        if result is True:
            processed += 1
        elif result is False:
            errors += 1
        else:
            skipped += 1

    print("\nSummary:")
    print(f"Processed: {processed}")
    print(f"Skipped: {skipped}")
    print(f"Errors: {errors}")


def main():
    parser = argparse.ArgumentParser(
        description="Generate instrumental tracks from audio stems"
    )
    parser.add_argument("root_folder", help="Root folder containing song subfolders")
    parser.add_argument("--song", help="Process only a specific song folder")

    args = parser.parse_args()

    if args.song:
        # Process single song
        song_path = Path(args.root_folder) / args.song
        if song_path.exists() and song_path.is_dir():
            create_instrumental(song_path)
        else:
            print(f"Song folder {song_path} not found")
    else:
        # Process all songs
        process_music_database(args.root_folder)


if __name__ == "__main__":
    main()