from flask import Flask, request
import json
import base64
from faster_whisper import WhisperModel
import tempfile
import os

def base64_to_wav(base64_data, save_path):
    wav_data = base64.b64decode(base64_data)
    with open(save_path, 'wb') as file:
        file.write(wav_data)
        
app = Flask(__name__)

GPU = True
wspr = WhisperModel("small", device="cuda" if GPU else "cpu", compute_type="int8")

@app.route('/api', methods=['GET' ,'POST'])
def STT():
    if request.method == 'POST':
        result = {}
        audio_data = request.data
        
        # Create a unique filename in the temporary directory
        temp_dir = tempfile.gettempdir()
        temp_file = tempfile.NamedTemporaryFile(suffix='.mp3', dir=temp_dir, delete=False)
        save_path = temp_file.name
        # save_path = "temp.wav"
        
#         save_path = "temp.wav"
        base64_to_wav(audio_data, save_path)

        segments, info = wspr.transcribe(save_path, beam_size=5, language="fr")
        texts = [el.text.strip() for el in segments]
        transcription = texts[0]
        query = transcription
        result["transcription"] = query
        if os.path.exists(save_path):
            os.remove(save_path)
        return result
    elif request.method == 'GET':
        # Add your test logic here
        return "API is working correctly!"

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5566, debug=True)