from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel

import asyncio
import tempfile
import subprocess
import os
import gc

import torch

app = FastAPI()

# Chỉ cho chạy 1 request mỗi lần
lock = asyncio.Semaphore(1)

# =========================
# Request model
# =========================

class TTSRequest(BaseModel):
    text: str

# =========================
# Root
# =========================

@app.get("/")
async def root():
    return {"status": "running"}

# =========================
# TTS Endpoint
# =========================

@app.post("/tts")
async def tts(req: TTSRequest):

    async with lock:

        # Kiểm tra text
        if not req.text.strip():
            return JSONResponse(
                status_code=400,
                content={
                    "success": False,
                    "error": "Text is empty"
                }
            )

        # Tạo file wav tạm
        temp_wav = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False
        )

        output_path = temp_wav.name
        temp_wav.close()

        try:

            # Command OmniVoice
            cmd = [
                "omnivoice-infer",
                "--text", req.text,
                "--output", output_path,
            ]

            print("Running command:")
            print(" ".join(cmd))

            # Chạy subprocess
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True
            )

            # Logs debug
            print("STDOUT:")
            print(result.stdout)

            print("STDERR:")
            print(result.stderr)

            # Nếu command lỗi
            if result.returncode != 0:
                return JSONResponse(
                    status_code=500,
                    content={
                        "success": False,
                        "error": result.stderr
                    }
                )

            # Không tạo file
            if not os.path.exists(output_path):
                return JSONResponse(
                    status_code=500,
                    content={
                        "success": False,
                        "error": "Audio file not created"
                    }
                )

            # File rỗng
            if os.path.getsize(output_path) == 0:
                return JSONResponse(
                    status_code=500,
                    content={
                        "success": False,
                        "error": "Audio file empty"
                    }
                )

            # Trả audio
            return FileResponse(
                path=output_path,
                media_type="audio/wav",
                filename="tts.wav"
            )

        except Exception as e:

            print("EXCEPTION:")
            print(str(e))

            return JSONResponse(
                status_code=500,
                content={
                    "success": False,
                    "error": str(e)
                }
            )

        finally:

            # Cleanup RAM
            gc.collect()

            if torch.cuda.is_available():
                torch.cuda.empty_cache()