from fastapi import FastAPI, UploadFile, File, Header, HTTPException
from fastapi.responses import JSONResponse
from transformers import pipeline
import torch
import io
import soundfile as sf
import os

app = FastAPI()

# Cấu hình Token bảo mật (Thay 'ChuyenGiaIoT2026' bằng mã bí mật của bạn)
API_TOKEN = os.getenv("ESP32_SECRET_TOKEN", "ChuyenGiaIoT2026")

print("Loading PhoWhisper model...")
# Sử dụng pipeline của Hugging Face để load mô hình PhoWhisper-small (tối ưu cho CPU Free)
stt_pipeline = pipeline(
    "automatic-speech-recognition",
    model="vinai/phowhisper-small",
    device="cpu"
)
print("Model loaded successfully!")

@app.get("/")
def read_root():
    return {"message": "Hugging Face STT Server cho ESP32-S3 đang hoạt động!"}

@app.post("/predict")
async def predict_speech(
    file: UploadFile = File(...),
    authorization: str = Header(None)
):
    # 1. Kiểm tra mã bảo mật Token từ ESP32 gửi lên
    if not authorization or authorization != f"Bearer {API_TOKEN}":
        raise HTTPException(status_code=401, detail="Unauthorized - Sai Token")
        
    try:
        # 2. Đọc file audio từ request
        audio_bytes = await file.read()
        
        # 3. Giải mã file WAV sang mảng numpy phù hợp với mô hình (ép về 16kHz)
        data, samplerate = sf.read(io.BytesIO(audio_bytes))
        
        # Nếu audio là Stereo, chuyển về Mono bằng cách lấy trung bình cộng
        if len(data.shape) > 1:
            data = data.mean(axis=1)
            
        # 4. Đưa vào mô hình AI để nhận diện Tiếng Việt
        # chunk_length_s=30 giúp xử lý tốt nếu file dài
        result = stt_pipeline({"raw": data, "sampling_rate": samplerate}, generate_kwargs={"language": "vi"})
        
        # 5. Trả kết quả gọn nhẹ về cho ESP32-S3
        return JSONResponse(content={
            "status": "success",
            "text": result["text"].strip()
        })
        
    except Exception as e:
        return JSONResponse(status_code=500, content={
            "status": "error",
            "message": str(e)
        })