import tempfile
import os
import numpy as np
import time
import logging
import os
import sys

from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect
from faster_whisper import WhisperModel

def env_bool(name: str, default: bool = False) -> bool:
    return os.getenv(name, str(default)).lower() in ("1", "true", "yes", "on")

LOG_LEVEL = os.getenv("WHISPER_LOG_LEVEL", "INFO").upper()
VAD_ENABLED = env_bool("WHISPER_VAD", True)
MODEL_SIZE = str(os.getenv("MODEL_SIZE", "small").lower())
MIN_AUDIO_SEC = float(os.getenv("MIN_AUDIO_SEC", "2.0"))
MAX_AUDIO_SEC = float(os.getenv("MAX_AUDIO_SEC", "10.0"))

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.INFO),
    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
    stream=sys.stdout,
    force=True
)

logger = logging.getLogger("whisper-ws")
logger.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))

logger.info(f"Log level = {LOG_LEVEL}")
logger.info(f"VAD enabled = {VAD_ENABLED}")

app = FastAPI()

# =========================
# Whisper model (shared)
# =========================
model = WhisperModel(MODEL_SIZE, device="cuda", compute_type="float32")

# =========================
# 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

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

    logger.info("WebSocket connecté !")

    audio_buffer = []
    start_time = time.time()

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

            logger.debug(f"Audio reçu: {len(data)} bytes")

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

            # silence pur → on reset le buffer
            if np.abs(pcm).mean() < 0.001:
                audio_buffer.clear()
                continue

            audio_buffer.append(pcm)

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

            logger.debug(f"Buffer audio: {audio_len:.2f}s")

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

            # Transcription dès qu'on a assez d'audio
            if audio_len >= MIN_AUDIO_SEC:
                logger.info(f"Transcription ({audio_len:.2f}s)")

                audio = np.concatenate(audio_buffer)
                audio_buffer.clear()

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

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

                
                logger.info(f"Texte: '{text}'")

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

                    logger.info("Résultat envoyé")

    except WebSocketDisconnect:
        logger.info("WebSocket déconnecté")
    except Exception as e:
        logger.exception("Erreur de WebSocket")

        await ws.send_json({"error": str(e)})
        await ws.close()