from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
import io
import soundfile as sf
from kittentts import KittenTTS

app = FastAPI(
    title="KittenTTS API",
    description="Full KittenTTS API with all models and voices",
    version="1.0.0"
)

# =========================
# AVAILABLE MODELS
# =========================

AVAILABLE_MODELS = {
    "mini": {
        "id": "KittenML/kitten-tts-mini-0.8",
        "params": "80M",
        "size": "80MB",
        "description": "Highest quality model"
    },
    "micro": {
        "id": "KittenML/kitten-tts-micro-0.8",
        "params": "40M",
        "size": "41MB",
        "description": "Balanced performance"
    },
    "nano": {
        "id": "KittenML/kitten-tts-nano-0.8",
        "params": "15M",
        "size": "56MB",
        "description": "Lightweight model"
    },
    "nano-int8": {
        "id": "KittenML/kitten-tts-nano-0.8-int8",
        "params": "15M",
        "size": "25MB",
        "description": "Ultra lightweight INT8 model"
    }
}

# =========================
# LOAD DEFAULT MODEL
# =========================

current_model_name = "nano-int8"
tts_model = KittenTTS(AVAILABLE_MODELS[current_model_name]["id"])

AVAILABLE_VOICES = [
    "Bella",
    "Jasper",
    "Luna",
    "Bruno",
    "Rosie",
    "Hugo",
    "Kiki",
    "Leo"
]

# =========================
# REQUEST MODEL
# =========================

class TTSRequest(BaseModel):
    text: str
    voice: str = "Bella"
    model: str = current_model_name
    speed: float = 1.0

# =========================
# ROOT ENDPOINT
# =========================

@app.get("/")
def root():
    return {
        "status": "running",
        "engine": "KittenTTS",
        "current_model": current_model_name,
        "voices": AVAILABLE_VOICES,
        "models": list(AVAILABLE_MODELS.keys())
    }

# =========================
# GET MODELS
# =========================

@app.get("/models")
def get_models():
    return AVAILABLE_MODELS

# =========================
# GET VOICES
# =========================

@app.get("/voices")
def get_voices():
    return {
        "voices": AVAILABLE_VOICES,
        "count": len(AVAILABLE_VOICES)
    }

# =========================
# SWITCH MODEL
# =========================

@app.post("/switch-model")
def switch_model(model_name: str):

    global tts_model, current_model_name

    if model_name not in AVAILABLE_MODELS:
        raise HTTPException(
            status_code=400,
            detail="Invalid model name"
        )

    tts_model = KittenTTS(
        AVAILABLE_MODELS[model_name]["id"]
    )

    current_model_name = model_name

    return {
        "status": "success",
        "current_model": model_name
    }

# =========================
# GENERATE SPEECH
# =========================

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

    global tts_model, current_model_name

    # switch model automatically if requested
    if req.model != current_model_name:

        if req.model not in AVAILABLE_MODELS:
            raise HTTPException(
                status_code=400,
                detail="Invalid model"
            )

        tts_model = KittenTTS(
            AVAILABLE_MODELS[req.model]["id"]
        )

        current_model_name = req.model

    # validate voice
    if req.voice not in AVAILABLE_VOICES:
        raise HTTPException(
            status_code=400,
            detail="Invalid voice"
        )

    try:

        audio = tts_model.generate(
            req.text,
            voice=req.voice
        )

        buffer = io.BytesIO()

        sf.write(
            buffer,
            audio,
            24000,
            format="WAV"
        )

        buffer.seek(0)

        return StreamingResponse(
            buffer,
            media_type="audio/wav",
            headers={
                "Content-Disposition":
                f"attachment; filename=kitten-{req.voice}.wav"
            }
        )

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )

# =========================
# API INFO
# =========================

@app.get("/info")
def info():

    return {
        "engine": "KittenTTS",
        "version": "0.8",
        "models_available": len(AVAILABLE_MODELS),
        "voices_available": len(AVAILABLE_VOICES),
        "current_model": current_model_name,
        "voices": AVAILABLE_VOICES,
        "models": AVAILABLE_MODELS
    }