import os
import sys
import soundfile as sf
from pedalboard import Pedalboard, Reverb, Compressor, HighShelfFilter, LowShelfFilter, Chorus, Distortion
from pedalboard.io import AudioFile
import numpy as np
from pathlib import Path

def apply_transformations_gpu(input_folder, output_folder):
    print(f"\n[STAGE 3] Applying Multi-Layer Transformations...")
    print(f"Reading from: {input_folder}")
    os.makedirs(output_folder, exist_ok=True)

    vocal_board = Pedalboard([Compressor(threshold_db=-20, ratio=4), Reverb(room_size=0.1, wet_level=0.1), HighShelfFilter(cutoff_frequency_hz=3000, gain_db=2)])
    drum_board = Pedalboard([Distortion(drive_db=2), LowShelfFilter(cutoff_frequency_hz=150, gain_db=3)])
    music_board = Pedalboard([Chorus(depth=0.2, rate_hz=1.0), Reverb(room_size=0.2, wet_level=0.15)])

    files_processed = 0
    # Walk through input folder to find stems
    for root, dirs, files in os.walk(input_folder):
        for file in files:
            if file.endswith(('.wav', '.mp3', '.flac')):
                input_path = os.path.join(root, file)
                output_path = os.path.join(output_folder, f"transformed_{file}")
                
                print(f"  Processing: {file}")
                with AudioFile(input_path) as f:
                    audio = f.read(f.frames)
                    samplerate = f.samplerate
                    board = vocal_board if "vocals" in file.lower() else (drum_board if "drums" in file.lower() else music_board)
                    effected = board(audio, samplerate)
                    with AudioFile(output_path, 'w', samplerate, effected.shape[0]) as o:
                        o.write(effected)
                files_processed += 1

    print(f"[SUCCESS] {files_processed} stems transformed and saved to: {output_folder}")

if __name__ == "__main__":
    stems_in = os.path.join(os.getcwd(), "Outputs", "Stems")
    fx_out = os.path.join(os.getcwd(), "Outputs", "Transformed")
    apply_transformations_gpu(stems_in, fx_out)