from flask import Flask, request, send_file
import soundfile as sf
import os
import requests
import subprocess
from pathlib import Path
import separate

app = Flask(__name__)

@app.route('/', methods=['GET'])
def hello():
    return "Hello! This is an api server, and it is running successfully. For usage, please contact the person who hosted this api server."

@app.route('/api/audio_separation', methods=['GET'])
def audio_separation():
    try:
    	# 获取MP3文件的直链地址
        mp3_url = request.args.get('url')
        if not mp3_url:
        	return "Error: URL parameter is required", 400

    	# 下载MP3文件到本地
        response = requests.get(mp3_url)
        mp3_filename = mp3_url.split('/')[-1]  # 使用下载的文件名
        with open("/tmp/" + mp3_filename, 'wb') as f:
        	f.write(response.content)

        # 执行音频分离操作
        audio_worker = separate.Predictor(args={
    											"files": [f"/tmp/{mp3_filename}"],
    											"output": Path("/tmp"),
    											"model_path": Path("./models/MDX_Net_Models/UVR-MDX-NET-Inst_HQ_3.onnx"),
    											"denoise": False,
    											"margin": 44100,
    											"chunks": 15,
    											"n_fft": 6144,
    											"dim_t": 8,
    											"dim_f": 2048
												})
        vocals, no_vocals, sampling_rate = audio_worker.predict("/tmp/" + mp3_filename)
        sf.write(os.path.join("/tmp", mp3_filename + "_no_vocals.wav"), no_vocals, sampling_rate)
        sf.write(os.path.join("/tmp", mp3_filename + "_vocals.wav"), vocals, sampling_rate)
        # subprocess.run(['python', 'separate.py', '/tmp/' + mp3_filename, '-o', '/tmp', '-m', './models/MDX_Net_Models/UVR-MDX-NET-Inst_HQ_3.onnx'])

    	# 生成分离后的文件名
        vocals_filename = f"{os.path.splitext(mp3_filename)[0]}_vocals.wav"
        no_vocals_filename = f"{os.path.splitext(mp3_filename)[0]}_no_vocals.wav"

    	# 提供文件的永久直链
        vocals_url = f"/download/{os.path.basename(vocals_filename)}"
        no_vocals_url = f"/download/{os.path.basename(no_vocals_filename)}"

    	# 返回结果
        result = {
        	"vocals_url": vocals_url,
        	"no_vocals_url": no_vocals_url
    	}
        return result
    except Exception as e:
        return "Error: " + str(e), 500

@app.route('/download/<filename>', methods=['GET'])
def download(filename):
    return send_file("/tmp/" + filename, as_attachment=True)

if __name__ == '__main__':
    app.run(debug=False)