from flask import Flask, request, jsonify
import os
import uuid
import subprocess
from werkzeug.utils import secure_filename
import time
import random

app = Flask(__name__)
UPLOAD_FOLDER = '/data'
ALLOWED_EXTENSIONS = {'wav', 'mp3', 'ogg', 'flac'}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def generate_filename():
    timestamp = int(time.time())
    random_num = random.randint(1000, 9999)
    return f"{timestamp}-{random_num}"

@app.route('/stt', methods=['POST'])
def transcribe():
    if 'audio' not in request.files:
        return jsonify({"error": "No audio file provided"}), 400

    file = request.files['audio']
    if file.filename == '':
        return jsonify({"error": "Empty filename"}), 400

    if not file or not allowed_file(file.filename):
        return jsonify({"error": "Invalid file type"}), 400

    try:
        # Generate filename with timestamp + random number
        file_base = generate_filename()
        orig_ext = secure_filename(file.filename).rsplit('.', 1)[1].lower()
        orig_path = os.path.join(UPLOAD_FOLDER, f"{file_base}.{orig_ext}")
        wav_path = os.path.join(UPLOAD_FOLDER, f"{file_base}.wav")
        output_base = os.path.join(UPLOAD_FOLDER, file_base)

        # Save original file
        file.save(orig_path)

        # Convert to WAV if needed
        if orig_ext != 'wav':
            subprocess.run([
                'ffmpeg', '-i', orig_path,
                '-ar', '16000', '-ac', '1', '-c:a', 'pcm_s16le', wav_path
            ], check=True)
            os.remove(orig_path)
            audio_path = wav_path
        else:
            audio_path = orig_path

        # Run whisper.cpp
        subprocess.run([
            '/app/whisper.cpp/build/bin/main',
            '-m', '/app/whisper.cpp/models/ggml-tiny.en-q5_1.bin',
            '-f', audio_path,
            '-otxt', '-of', output_base
        ], check=True)

        # Read output
        output_file = f"{output_base}.txt"
        with open(output_file, 'r') as f:
            transcription = f.read()

        # Return the filenames in the response
        response = {
            "text": transcription.strip(),
            "audio_file": f"{file_base}.{orig_ext}",
            "text_file": f"{file_base}.txt"
        }

        return jsonify(response)

    except subprocess.CalledProcessError as e:
        return jsonify({"error": f"Whisper processing failed: {e.stderr}"}), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    os.makedirs(UPLOAD_FOLDER, exist_ok=True)
    app.run(host='0.0.0.0', port=4004)