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


def detect_vocal_boundaries(
    audio_path,
    sr=44100,
    frame_length=2048,
    hop_length=512,
    silence_threshold=-40,
    min_silence_duration=0.5,
):
    """
    Detect the start and end of vocal content in an audio file.

    Args:
        audio_path: Path to the audio file
        sr: Sample rate for loading audio
        frame_length: Frame length for RMS calculation
        hop_length: Hop length for RMS calculation
        silence_threshold: Threshold in dB below which audio is considered silence
        min_silence_duration: Minimum duration of silence (in seconds) to consider for trimming

    Returns:
        tuple: (start_time, end_time) in seconds
    """
    # Load audio
    y, original_sr = librosa.load(audio_path, sr=sr)

    # Calculate RMS energy
    rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]

    # Convert to dB
    rms_db = librosa.amplitude_to_db(rms, ref=np.max)

    # Find frames above threshold
    active_frames = rms_db > silence_threshold

    # Convert frame indices to time
    times = librosa.frames_to_time(np.arange(len(rms_db)), sr=sr, hop_length=hop_length)

    # Find first and last active frames
    active_indices = np.where(active_frames)[0]

    if len(active_indices) == 0:
        # If no active frames found, return full duration
        return 0, librosa.get_duration(y=y, sr=sr)

    start_frame = active_indices[0]
    end_frame = active_indices[-1]

    # Convert to time, with some padding to avoid cutting off vocals
    padding = 0.1  # 100ms padding
    start_time = max(0, times[start_frame] - padding)
    end_time = min(librosa.get_duration(y=y, sr=sr), times[end_frame] + padding)

    # Ensure minimum duration
    if end_time - start_time < min_silence_duration:
        return 0, librosa.get_duration(y=y, sr=sr)

    return start_time, end_time


def trim_audio_file(
    input_path, output_path, start_time, end_time, max_duration=30.0, sr=44100
):
    """
    Trim an audio file between start_time and end_time, then limit to max_duration.

    Args:
        input_path: Path to input audio file
        output_path: Path to save trimmed audio
        start_time: Start time in seconds
        end_time: End time in seconds
        max_duration: Maximum duration of output in seconds (default 30.0)
        sr: Sample rate for processing
    """
    # Load audio
    y, original_sr = librosa.load(input_path, sr=None)  # Keep original sample rate

    # Convert times to sample indices
    start_sample = int(start_time * original_sr)
    end_sample = int(end_time * original_sr)

    # First trim to vocal boundaries
    y_trimmed = y[start_sample:end_sample]

    # Then limit to max_duration (30 seconds) from the start of trimmed audio
    max_samples = int(max_duration * original_sr)
    if len(y_trimmed) > max_samples:
        y_trimmed = y_trimmed[:max_samples]

    # Save trimmed audio
    sf.write(output_path, y_trimmed, original_sr)


def process_songs(parent_folder, reference_filename="vocals"):
    """
    Process all songs in the dataset, trimming based on specified reference file.

    Args:
        parent_folder: Path to the parent folder containing all dataset folders
        reference_filename: Name of the audio file to use for silence detection (without .wav extension)
    """
    parent_path = Path(parent_folder)

    # Define folder names
    folders = ["original_format", "suno_format", "lalal_format"]

    # Create output folders with reference filename in the name
    for folder in folders:
        output_folder = parent_path / f"{folder}_trimmed_{reference_filename}"
        output_folder.mkdir(exist_ok=True)

    # Get ground truth folder
    ground_truth_path = parent_path / "original_format"

    if not ground_truth_path.exists():
        print(f"Ground truth folder not found: {ground_truth_path}")
        return

    # Get all song folders from ground truth
    song_folders = [f for f in ground_truth_path.iterdir() if f.is_dir()]

    print(f"Found {len(song_folders)} songs to process")
    print(f"Using '{reference_filename}.wav' as reference for silence detection")

    # Process each song
    for song_folder in song_folders:
        song_name = song_folder.name
        print(f"\nProcessing song: {song_name}")

        # Path to reference audio file for silence detection
        reference_audio_path = song_folder / f"{reference_filename}.wav"

        if not reference_audio_path.exists():
            print(
                f"  Warning: Reference file '{reference_filename}.wav' not found for {song_name}"
            )
            continue

        try:
            # Detect boundaries from reference file
            start_time, end_time = detect_vocal_boundaries(str(reference_audio_path))
            duration_after_trim = end_time - start_time
            final_duration = min(duration_after_trim, 30.0)
            print(
                f"  Detected active content in {reference_filename}: {start_time:.2f}s to {end_time:.2f}s"
            )
            print(f"  Final output duration: {final_duration:.2f}s")

            # Process all folders for this song
            for folder in folders:
                source_folder = parent_path / folder / song_name
                output_folder = (
                    parent_path / f"{folder}_trimmed_{reference_filename}" / song_name
                )

                if not source_folder.exists():
                    print(f"  Warning: Source folder not found: {source_folder}")
                    continue

                # Create output song folder
                output_folder.mkdir(parents=True, exist_ok=True)

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

                for wav_file in wav_files:
                    file_stem = wav_file.stem  # filename without extension

                    if file_stem == "mixture":
                        # Special handling for mixture - rename to indicate reference
                        output_filename = f"mixture_{reference_filename}_trimmed.wav"
                    else:
                        # Regular files keep their original names
                        output_filename = f"{file_stem}.wav"

                    output_path = output_folder / output_filename

                    try:
                        trim_audio_file(
                            str(wav_file), str(output_path), start_time, end_time
                        )
                        print(
                            f"    Trimmed {file_stem}: {folder}/{song_name}/{output_filename}"
                        )
                    except Exception as e:
                        print(f"    Error trimming {file_stem}: {str(e)}")

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

    print(
        f"\nProcessing complete! Trimmed files saved in folders with '_trimmed_{reference_filename}' suffix."
    )


def main():
    parser = argparse.ArgumentParser(
        description="Audio Trimming Script for Silence Removal"
    )
    parser.add_argument(
        "--parent_folder",
        type=str,
        default="/app2/suno/data/sara/musdb",
        help="Path to the parent folder containing all dataset folders",
    )
    parser.add_argument(
        "--reference_file",
        type=str,
        default="vocals",
        help="Name of the audio file to use for silence detection (without .wav extension). Examples: vocals, bass, drums, piano",
    )

    args = parser.parse_args()

    print("Enhanced Audio Trimming Script for Silence Removal")
    print("=" * 55)
    print(f"Processing folder: {args.parent_folder}")
    print(f"Reference file: {args.reference_file}.wav")

    if not os.path.exists(args.parent_folder):
        print(f"Error: Parent folder does not exist: {args.parent_folder}")
        return

    # Process all songs
    process_songs(args.parent_folder, args.reference_file)


if __name__ == "__main__":
    main()
