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é !")

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

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

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

            energy = float(np.abs(pcm).mean())
            logger.debug(f"Énergie audio = {energy:.6f}")

            # silence gate léger
            if np.abs(pcm).mean() < 0.00005:
                logger.debug("Chunk silencieux ignoré")
                continue

            duration = len(pcm) / SAMPLE_RATE
            logger.info(f"Transcription chunk {duration:.2f}s")

            segments, _ = model.transcribe(
                pcm,
                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
                })

    except WebSocketDisconnect:
        logger.info("WebSocket déconnecté")
    except Exception:
        logger.exception("Erreur WebSocket")
        await ws.close()