from flask import Flask, request, send_file, render_template
import subprocess
import os
import sys
from datetime import datetime
from io import BytesIO

app = Flask(__name__)

def run_f5_tts(ref_audio_path, ref_text, gen_text, model="F5TTS_Base", speed=1.2, vocoder_name="vocos"):
    current_dir = os.path.dirname(os.path.abspath(__file__))
    infer_cli_path = os.path.join(current_dir, "src", "f5_tts", "infer", "infer_cli.py")
    vocab_file = os.path.join(current_dir, "F5-TTS-MRSU", "vocab.txt")
    ckpt_file = os.path.join(current_dir, "F5-TTS-MRSU", "model_last.pt")
    tests_dir = os.path.join(current_dir, "tests")

    os.environ['PYTHONIOENCODING'] = 'utf-8'

    # Tăng speed để thử nghiệm (tùy chỉnh theo tài liệu)
    command = [
        sys.executable,
        infer_cli_path,
        "--model", model,
        "--ref_audio", ref_audio_path,
        "--ref_text", ref_text,
        "--gen_text", gen_text,
        "--speed", str(speed),  # Tăng tốc độ lên 1.2
        "--vocoder_name", vocoder_name,
        "--vocab_file", vocab_file,
        "--ckpt_file", ckpt_file
    ]

    try:
        result = subprocess.run(
            command,
            check=True,
            capture_output=True,
            text=True,
            encoding='utf-8'
        )
        if os.path.exists(tests_dir):
            wav_files = [f for f in os.listdir(tests_dir) if f.endswith('.wav')]
            if wav_files:
                latest_wav = max(wav_files, key=lambda x: os.path.getmtime(os.path.join(tests_dir, x)))
                output_file = os.path.join(tests_dir, latest_wav)
                return True, output_file
        return False, "Không tìm thấy file âm thanh trong thư mục tests"
    except subprocess.CalledProcessError as e:
        return False, e.stderr
    except Exception as e:
        return False, str(e)

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

@app.route('/api/generate', methods=['POST'])
def generate_speech():
    ref_audio = request.files['ref_audio']
    ref_text = request.form['ref_text']
    gen_text = request.form['gen_text']
    model = request.form.get('model', 'F5TTS_Base')
    speed = float(request.form.get('speed', 1.2))  # Mặc định tăng tốc độ

    ref_audio_path = 'temp_ref.wav'
    ref_audio.save(ref_audio_path)

    success, result = run_f5_tts(ref_audio_path, ref_text, gen_text, model, speed)
    os.remove(ref_audio_path)

    if success:
        return send_file(result, mimetype='audio/wav')
    else:
        return result, 500

if __name__ == "__main__":
    app.run(debug=True)