import os
import io
import base64
import torch
import soundfile as sf
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModel, AutoTokenizer

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_ID = "ai4bharat/vits_rasa_13"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Preload model weights into memory
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
model.eval()

class TTSRequest(BaseModel):
    text: str
    gender: str = "male"  # male = 5 (BRX_M), female = 4 (BRX_F)

@app.get("/")
def health():
    return {"status": "running", "engine": "ai4bharat-bodo-vits"}

@app.post("/synthesize")
def synthesize(req: TTSRequest):
    text = req.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="Text cannot be empty.")

    # Speaker mapping: 5 = Bodo Male, 4 = Bodo Female
    # Style mapping: 4 = Conversational (CONV)
    speaker_id = 5 if req.gender == "male" else 4
    style_id = 4

    try:
        inputs = tokenizer(text=text, return_tensors="pt")
        with torch.no_grad():
            outputs = model(inputs["input_ids"], speaker_id=speaker_id, emotion_id=style_id)

        buffer = io.BytesIO()
        sf.write(buffer, outputs.waveform.squeeze().cpu().numpy(), model.config.sampling_rate, format="WAV")
        b64_audio = base64.b64encode(buffer.getvalue()).decode("utf-8")

        return {"audioUrl": f"data:audio/wav;base64,{b64_audio}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))