from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from kokoro import KPipeline
import io
import soundfile as sf
import os
from huggingface_hub import InferenceClient

# Retrieve the token securely from the environment
hf_token = os.getenv("HF_TOKEN")

# Example usage
client = InferenceClient(token=hf_token)

app = FastAPI(title="Kokoro TTS API Server")

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

# Initialize the pipeline into memory once at boot
pipeline = KPipeline(lang_code='a')

# FIXED: Listening to the clean root path so Hugging Face forwards traffic directly!
@app.post("/")
async def text_to_speech(payload: dict):
    text = payload.get("input", "")
    voice = payload.get("voice", "af_sky")
    
    if not text:
        raise HTTPException(status_code=400, detail="Missing text input")
        
    def audio_stream_generator():
        try:
            generator = pipeline(text, voice=voice, speed=1, split_pattern=r'\n+')
            for _, _, audio in generator:
                byte_io = io.BytesIO()
                sf.write(byte_io, audio, 24000, format='MP3')
                yield byte_io.getvalue()
        except Exception as e:
            print(f"Streaming error: {e}")

    return StreamingResponse(audio_stream_generator(), media_type="audio/mpeg")