import os
import wave
import json
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS
from vosk import Model, KaldiRecognizer
from flasgger import Swagger

# Thư mục chứa model
MODEL_PATH = "model/vosk-model"

# Kiểm tra model đã tải chưa
if not os.path.exists(MODEL_PATH):
    raise Exception(f"❌ Model Vosk không tìm thấy tại {MODEL_PATH}! Kiểm tra lại.")

# Load model
print("✅ Đang tải model Vosk...")
model = Model(MODEL_PATH)

# Khởi tạo Flask app
app = Flask(__name__)
CORS(app)
swagger = Swagger(app)

@app.route("/")
def home():
    """API Home
    ---
    responses:
      200:
        description: API đang chạy
    """
    return "✅ Vosk STT API đang chạy!"

@app.route("/stt", methods=["POST"])
def stt():
    """Chuyển đổi giọng nói thành văn bản (Speech-to-Text)
    ---
    consumes:
      - multipart/form-data
    parameters:
      - in: formData
        name: audio
        type: file
        required: true
        description: File âm thanh WAV mono PCM
    responses:
      200:
        description: Kết quả chuyển đổi văn bản
        schema:
          type: object
          properties:
            text:
              type: string
              example: "Xin chào thế giới"
      400:
        description: Lỗi nếu file âm thanh không hợp lệ hoặc không tìm thấy
    """
    if "audio" not in request.files:
        return jsonify({"error": "Không tìm thấy file audio!"}), 400

    audio_file = request.files["audio"]
    file_path = "temp.wav"
    audio_file.save(file_path)

    # Mở file âm thanh
    wf = wave.open(file_path, "rb")

    # Kiểm tra file có đúng định dạng WAV mono không
    if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getcomptype() != "NONE":
        os.remove(file_path)
        return jsonify({"error": "File audio phải là WAV mono PCM!"}), 400

    rec = KaldiRecognizer(model, wf.getframerate())
    result_text = ""

    while True:
        data = wf.readframes(4000)
        if len(data) == 0:
            break
        if rec.AcceptWaveform(data):
            result_text += json.loads(rec.Result())["text"] + " "

    wf.close()
    os.remove(file_path)  # Xóa file tạm

    return jsonify({"text": result_text.strip()})

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