import os
import io
import soundfile as sf
from flask import Flask, request, Response, stream_with_context, jsonify
from flask_cors import CORS
from kokoro_onnx import Kokoro

app = Flask(__name__)
CORS(app) # Enable CORS for SillyTavern

# Initialize (0.4.8 works automatically with Docker's espeak)
print("--- LOADING KOKORO v0.4.8 ---")
try:
    # Note: 0.4.8 auto-detects voices.bin/json
    kokoro = Kokoro("kokoro-v1.0.int8.onnx", "voices-v1.0.bin")
    print("SUCCESS: Model Loaded!")
except Exception as e:
    print(f"FATAL ERROR: {e}")
    kokoro = None

@app.route("/", methods=["GET"])
def home():
    return "<h1>Kokoro API is Online (v0.4.8 Stable)</h1>"

# Universal Route (Accepts both /v1 and standard paths)
@app.route("/audio/speech", methods=["POST", "OPTIONS"])
@app.route("/v1/audio/speech", methods=["POST", "OPTIONS"])
def generate_speech():
    if request.method == "OPTIONS":
        return Response(status=200)

    if not kokoro:
        return jsonify({"error": "Server crashed at startup - check logs"}), 500

    try:
        data = request.json
        text = data.get("input", "")
        voice = data.get("voice", "af_bella")
        speed = float(data.get("speed", 1.0))

        def generate_stream():
            samples, sample_rate = kokoro.create(
                text, 
                voice=voice, 
                speed=speed, 
                lang="en-us"
            )
            buffer = io.BytesIO()
            sf.write(buffer, samples, sample_rate, format='WAV')
            buffer.seek(0)
            yield buffer.read()

        return Response(stream_with_context(generate_stream()), mimetype="audio/wav")
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/v1/models", methods=["GET", "OPTIONS"])
def list_models():
    return jsonify({"data": [{"id": "kokoro", "object": "model"}]})

@app.route("/v1/audio/voices", methods=["GET", "OPTIONS"])
def list_voices():
    if kokoro:
        return jsonify({"voices": list(kokoro.get_voices())})
    return jsonify({"voices": []})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)