import os
import wave
import json
import asyncio
import websockets
import uuid
from vosk import Model, KaldiRecognizer

# 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"\u274C Model Vosk không tìm thấy tại {MODEL_PATH}! Kiểm tra lại.")

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

async def recognize_audio(websocket, path):
    print("🔵 Kết nối WebSocket mới...")
    rec = KaldiRecognizer(model, 16000)
    result_text = ""
    
    try:
        async for message in websocket:
            if isinstance(message, bytes):  # Kiểm tra nếu dữ liệu là bytes
                if rec.AcceptWaveform(message):
                    result = json.loads(rec.Result())
                    text = result.get("text", "")
                    result_text += text + " "
                    await websocket.send(json.dumps({"text": text}))
            else:
                await websocket.send(json.dumps({"error": "Dữ liệu không hợp lệ"}))
    except websockets.exceptions.ConnectionClosed:
        print("🔴 Kết nối WebSocket đã đóng")
    finally:
        print(f"✅ Văn bản cuối cùng: {result_text.strip()}")

# Chạy server WebSocket
start_server = websockets.serve(recognize_audio, "0.0.0.0", 7860)
print("🚀 WebSocket STT server đang chạy trên ws://0.0.0.0:7860")

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()