from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, JSONResponse
import outetts
import io
import json
import asyncio
import os

# Initialize the interface
interface = outetts.Interface(
    config=outetts.ModelConfig.auto_config(
        model=outetts.Models.VERSION_1_0_SIZE_1B,
        # For llama.cpp backend
        # backend=outetts.Backend.LLAMACPP,
        # quantization=outetts.LlamaCppQuantization.FP16
        # For transformers backend
        backend=outetts.Backend.HF,
    )
)

# Load the default speaker profile
speaker = interface.load_default_speaker("EN-FEMALE-1-NEUTRAL")

app = FastAPI()

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

async def process_chunk(text_chunk: str, websocket: WebSocket):
    try:
        output = interface.generate(
            config=outetts.GenerationConfig(
                text=text_chunk,
                generation_type=outetts.GenerationType.CHUNKED,
                speaker=speaker,
                sampler_config=outetts.SamplerConfig(
                    temperature=0.4
                ),
            )
        )
        # Save audio to buffer
        audio_buffer = io.BytesIO()
        output.save(audio_buffer)
        audio_buffer.seek(0)
        audio_bytes = audio_buffer.read()
        # Send audio bytes back
        await websocket.send_bytes(audio_bytes)
    except Exception as e:
        await websocket.send_text(json.dumps({"error": str(e)}))

@app.websocket("/ws/tts")
async def websocket_tts(websocket: WebSocket):
    await websocket.accept()
    tasks: set[asyncio.Task] = set()
    try:
        while True:
            data = await websocket.receive_text()
            # Schedule processing without awaiting
            task = asyncio.create_task(process_chunk(data, websocket))
            tasks.add(task)
            task.add_done_callback(lambda t: tasks.discard(t))
    except WebSocketDisconnect:
        # Cancel all pending tasks
        for task in tasks:
            task.cancel()