import os
import tempfile
import subprocess
import torch
import whisper
import soundfile as sf
from fastapi import FastAPI, UploadFile
import uvicorn
import psutil
import platform

app = FastAPI()

# Choix du device : GPU si dispo, sinon CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model = whisper.load_model("medium", device=device)


def convert_to_wav(input_path: str, output_path: str):
    """
    Conversion de l'audio en WAV 16kHz mono PCM avec ffmpeg.
    """
    cmd = [
        "ffmpeg", "-y", "-i", input_path,
        "-ar", "16000", "-ac", "1", "-f", "wav", output_path
    ]
    subprocess.run(cmd, check=True)

@app.get("/status/")
async def status():
    # Infos système
    cpu_percent = psutil.cpu_percent(interval=0.5)
    memory = psutil.virtual_memory()
    gpu_info = None
    
    if torch.cuda.is_available():
        gpu_info = {
            "name": torch.cuda.get_device_name(0),
            "memory_total_GB": round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 2),
            "memory_reserved_GB": round(torch.cuda.memory_reserved(0) / (1024**3), 2),
            "memory_allocated_GB": round(torch.cuda.memory_allocated(0) / (1024**3), 2),
        }

    return {
        "status": "OK",
        "model": model,
        "device": device,
        "system": {
            "os": platform.system(),
            "os_version": platform.version(),
            "python_version": platform.python_version(),
        },
        "cpu": {
            "usage_percent": cpu_percent
        },
        "memory": {
            "total_GB": round(memory.total / (1024**3), 2),
            "available_GB": round(memory.available / (1024**3), 2),
            "used_percent": memory.percent
        },
        "gpu": gpu_info
    }
        
@app.post("/transcribe/")
async def transcribe(file: UploadFile):
    if not file.filename.lower().endswith((".wav", ".mp3")):
        return {"error": "Seuls les fichiers WAV ou MP3 sont acceptés"}

    # Sauvegarde du fichier uploadé
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp_in:
        audio_data = await file.read()
        tmp_in.write(audio_data)
        input_path = tmp_in.name

    # Conversion en WAV standard
    tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    tmp_out.close()
    convert_to_wav(input_path, tmp_out.name)

    # Log durée audio (pratique pour debug)
    f = sf.SoundFile(tmp_out.name)
    duration = len(f) / f.samplerate
    print(f"[INFO] Durée du fichier audio : {duration:.2f} secondes")

    try:
        # Transcription avec paramètres robustes
        result = model.transcribe(
            tmp_out.name,
            language="fr",
            fp16=False,
            temperature=0,
            word_timestamps=True
        )

        return {
            "text": result["text"],
            "segments": [
                {
                    "start": seg["start"],
                    "end": seg["end"],
                    "text": seg["text"]
                }
                for seg in result["segments"]
            ],
            "duration_seconds": duration,
            "language": result.get("language", "fr")
        }
    finally:
        if os.path.exists(input_path):
            os.remove(input_path)
        if os.path.exists(tmp_out.name):
            os.remove(tmp_out.name)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
