from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect
from faster_whisper import WhisperModel
import tempfile
import os
import numpy as np
import time

app = FastAPI()

# =========================
# Whisper model (shared)
# =========================
model = WhisperModel(
    "small",
    device="cuda",
    compute_type="float32"  # GTX 1050 Ti compatible
)

# =========================
# REST API (Swagger / tests)
# =========================
@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...)):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    segments, info = model.transcribe(
        tmp_path,
        language="fr",
        vad_filter=True
    )

    os.remove(tmp_path)

    text = " ".join(seg.text for seg in segments)

    return {
        "text": text,
        "language": info.language
    }

# =========================
# WebSocket STT (streaming)
# =========================

SAMPLE_RATE = 16000
MIN_AUDIO_SEC = 0.5
MAX_AUDIO_SEC = 10.0

@app.websocket("/ws/transcribe")
async def ws_transcribe(ws: WebSocket):
    await ws.accept()

    audio_buffer = []
    start_time = time.time()

    try:
        while True:
            data = await ws.receive_bytes()

            # PCM int16 → float32
            pcm = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
            audio_buffer.append(pcm)

            audio_len = sum(len(chunk) for chunk in audio_buffer) / SAMPLE_RATE

            # Sécurité : reset si trop long
            if audio_len > MAX_AUDIO_SEC:
                audio_buffer.clear()
                start_time = time.time()
                continue

            # Transcription dès qu'on a assez d'audio
            if audio_len >= MIN_AUDIO_SEC:
                audio = np.concatenate(audio_buffer)
                audio_buffer.clear()

                segments, _ = model.transcribe(
                    audio,
                    language="fr",
                    vad_filter=True
                )

                text = " ".join(seg.text for seg in segments).strip()

                if text:
                    await ws.send_json({
                        "type": "final",
                        "text": text
                    })

    except WebSocketDisconnect:
        pass
    except Exception as e:
        await ws.send_json({"error": str(e)})
        await ws.close()