import os
import librosa
import soundfile as sf
import noisereduce as nr
from flask import Flask, render_template, request, send_from_directory, jsonify
from flask_socketio import SocketIO, emit
from pedalboard import Pedalboard, Reverb, PitchShift
from pydub import AudioSegment
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['PROCESSED_FOLDER'] = 'processed'
socketio = SocketIO(app, cors_allowed_origins="*")

# Create folders if they don't exist
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['PROCESSED_FOLDER'], exist_ok=True)

# Shared variable for the latest processed track
processed_info = {'y': None, 'sr': None, 'filename': ''}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload():
    file = request.files['file']
    filename = secure_filename(file.filename)
    file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
    return jsonify({"filename": filename})

@socketio.on('start_processing')
def handle_audio(data):
    filename = data['filename']
    path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    
    # 1. Loading
    emit('progress', {'msg': 'Loading audio file...'})
    y, sr = librosa.load(path, sr=None)

    # 2. Denoise
    emit('progress', {'msg': 'Cleaning vocal noise...'})
    y_clean = nr.reduce_noise(y=y, sr=sr)

    # 3. Autotune
    emit('progress', {'msg': 'Applying pitch correction...'})
    board = Pedalboard([PitchShift(semitones=0)]) 
    y_tuned = board(y_clean, sr)

    # 4. BPM/Scale
    emit('progress', {'msg': 'Calculating BPM and Scale...'})
    tempo, _ = librosa.beat.beat_track(y=y_tuned, sr=sr)
    emit('progress', {'msg': f'Detected Tempo: {int(tempo)} BPM'})

    # 5. Reverb/Echo
    emit('progress', {'msg': 'Adding reverb and echo...'})
    fx = Pedalboard([Reverb(room_size=0.3, wet_level=0.2)])
    y_final = fx(y_tuned, sr)

    # 6. Finalize
    emit('progress', {'msg': 'Layering instruments... (Complete)'})
    global processed_info
    processed_info = {'y': y_final, 'sr': sr, 'name': filename.split('.')[0]}
    emit('progress', {'msg': 'Production finished!'})

@app.route('/download/<fmt>')
def download(fmt):
    global processed_info
    base = processed_info['name']
    wav_path = os.path.join(app.config['PROCESSED_FOLDER'], f"{base}.wav")
    sf.write(wav_path, processed_info['y'], processed_info['sr'])
    
    if fmt == 'mp3':
        mp3_path = os.path.join(app.config['PROCESSED_FOLDER'], f"{base}.mp3")
        AudioSegment.from_wav(wav_path).export(mp3_path, format="mp3")
        return send_from_directory(app.config['PROCESSED_FOLDER'], f"{base}.mp3")
    
    return send_from_directory(app.config['PROCESSED_FOLDER'], f"{base}.wav")

if __name__ == '__main__':
    socketio.run(app, debug=True)