import tempfile
import os
import numpy as np
import time
import logging
import os
import sys
import asyncio
from threading import Lock

from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect, HTTPException
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())
COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "float32").strip().lower()
BEAM_SIZE = int(os.getenv("WHISPER_BEAM_SIZE", "1"))
if BEAM_SIZE < 1:
    raise ValueError("WHISPER_BEAM_SIZE doit être supérieur ou égal à 1")
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}")
logger.info("Model = %s, compute type = %s", MODEL_SIZE, COMPUTE_TYPE)
logger.info("Beam size = %d", BEAM_SIZE)

app = FastAPI()

# =========================
# Whisper model (shared)
# =========================
model = WhisperModel(MODEL_SIZE, device="cuda", compute_type=COMPUTE_TYPE)
# Vérifier les kernels CUDA avant de déclarer le service prêt.
list(model.transcribe(np.zeros(16000, dtype=np.float32), language="fr", beam_size=BEAM_SIZE)[0])
model_lock = Lock()

def transcribe_audio(audio):
    # Consommer le générateur sous le verrou, hors de la boucle asyncio.
    with model_lock:
        segments, info = model.transcribe(
            audio, language="fr", vad_filter=VAD_ENABLED,
            beam_size=BEAM_SIZE, condition_on_previous_text=False,
        )
        return " ".join(seg.text for seg in segments).strip(), info.language

@app.get("/healthz")
def healthz():
    return {"status": "ready", "model": MODEL_SIZE, "compute_type": COMPUTE_TYPE, "beam_size": BEAM_SIZE}

# =========================
# REST API (Swagger / tests)
# =========================
@app.post("/transcribe")
def transcribe(file: UploadFile = File(...)):
    # Discord envoie au maximum 20 secondes de WAV stéréo 48 kHz.
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp_path = tmp.name
        try:
            data = file.file.read(8 * 1024 * 1024 + 1)
            if len(data) > 8 * 1024 * 1024:
                raise HTTPException(413, "Audio trop volumineux")
            tmp.write(data)
            tmp.flush()
            text, language = transcribe_audio(tmp_path)
            return {"text": text, "language": language}
        finally:
            os.remove(tmp_path)

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

            text, _ = await asyncio.to_thread(transcribe_audio, pcm)

            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()