from flask import Flask, request, send_file, jsonify
from pocket_tts import TTSModel
import scipy.io.wavfile
import tempfile
import torch

# =========================
# CPU Optimization
# =========================
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

app = Flask(__name__)

# =========================
# Hardcoded Voice
# =========================
VOICE = "eve" #"jane" #"vera" #"eve" #"anna"

print("Loading TTS model...")

# Load model
model = TTSModel.load_model()

print("Preparing voice...")

# Built-in voice
voice_state = model.get_state_for_audio_prompt(
    VOICE
)

print("TTS Server Ready!")

# =========================
# Home Route
# =========================
@app.route("/", methods=["GET"])
def home():
    return {
        "status": "running",
        "voice": VOICE,
        "usage": "/generate?q=hello"
    }

# =========================
# Generate Route
# =========================
@app.route("/generate", methods=["GET"])
def generate_get():

    text = request.args.get("q")

    if not text:
        return jsonify({
            "error": "Missing ?q=text"
        }), 400

    try:

        # Generate speech
        audio = model.generate_audio(
            voice_state,
            text
        )

        # Temp WAV file
        temp_wav = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False
        )

        # Save WAV
        scipy.io.wavfile.write(
            temp_wav.name,
            model.sample_rate,
            audio.numpy()
        )

        # Return audio
        return send_file(
            temp_wav.name,
            mimetype="audio/wav",
            as_attachment=False
        )

    except Exception as e:
        return jsonify({
            "error": str(e)
        }), 500


@app.route("/generate", methods=["POST"])
def generate():

    # Get JSON body
    data = request.get_json()

    # Get q from body
    text = data.get("q") if data else None

    if not text:
        return jsonify({
            "error": "Missing 'q' in JSON body"
        }), 400

    try:

        # Generate speech
        audio = model.generate_audio(
            voice_state,
            text
        )

        # Temp WAV file
        temp_wav = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False
        )

        # Save WAV
        scipy.io.wavfile.write(
            temp_wav.name,
            model.sample_rate,
            audio.numpy()
        )

        # Return audio
        return send_file(
            temp_wav.name,
            mimetype="audio/wav",
            as_attachment=False
        )

    except Exception as e:
        return jsonify({
            "error": str(e)
        }), 500

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