from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel
import uuid
import os
import json
import torch
from TTS.api import TTS

# ---------------------------
# CONFIG
# ---------------------------
ALLOWED_API_KEYS = ["your_master_key_here"]  # You can add/remove keys
MODEL_NAME = "coqui/XTTS-v2"
OUTPUT_DIR = "outputs"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ---------------------------
# LOAD MODEL (CPU mode)
# ---------------------------
device = "cpu"
tts = TTS(model_name=MODEL_NAME).to(device)

# ---------------------------
# FASTAPI INIT
# ---------------------------
app = FastAPI()


# ---------------------------
# REQUEST BODY
# ---------------------------
class TTSRequest(BaseModel):
    api_key: str
    text: str
    language: str = "en"
    speaker_wav: str | None = None


# ---------------------------
# ROOT
# ---------------------------
@app.get("/")
def root():
    return {"message": "XTTS Custom API Running Successfully!"}


# ---------------------------
# TTS ENDPOINT
# ---------------------------
@app.post("/generate")
def generate_audio(req: TTSRequest):

    # API KEY VALIDATION
    if req.api_key not in ALLOWED_API_KEYS:
        raise HTTPException(status_code=401, detail="Invalid API Key")

    # FILE NAME
    file_id = str(uuid.uuid4())
    out_file = f"{OUTPUT_DIR}/{file_id}.wav"

    # RUN TTS
    try:
        tts.tts_to_file(
            text=req.text,
            file_path=out_file,
            speaker_wav=req.speaker_wav,
            language=req.language
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    # RETURN AUDIO FILE
    return {"status": "success", "audio_url": f"/audio/{file_id}.wav"}


# ---------------------------
# AUDIO FILE SERVE
# ---------------------------
@app.get("/audio/{file_name}")
def get_audio(file_name: str):
    file_path = os.path.join(OUTPUT_DIR, file_name)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(path=file_path, media_type="audio/wav")