import uuid
import threading
from flask import Flask, jsonify, send_file
from pocket_tts_onnx import PocketTTSOnnx
import soundfile as sf
import numpy as np
import os

app = Flask(__name__)
tts = PocketTTSOnnx(precision="int8")

# In-memory store for task status (In production, use Redis)
tasks = {}

def generate_audio_task(task_id, text, voice_path):
    """Background task to generate full audio."""
    # Generate the full audio array
    full_audio = tts.predict(text, voice=voice_path)
    
    # Save to a temporary file
    filename = f"output_{task_id}.wav"
    sf.write(filename, full_audio, 24000)
    
    tasks[task_id] = {"status": "completed", "file": filename}

@app.route("/generate", methods=["POST"])
def start_generation():
    text = request.json.get("text")
    voice = request.json.get("voice", "samples/reference.wav")
    task_id = str(uuid.uuid4())
    
    tasks[task_id] = {"status": "processing"}
    
    # Start inference in a background thread
    threading.Thread(target=generate_audio_task, args=(task_id, text, voice)).start()
    
    return jsonify({"task_id": task_id, "message": "Generation started"})

@app.route("/status/<task_id>")
def check_status(task_id):
    return jsonify(tasks.get(task_id, {"status": "not_found"}))

@app.route("/download/<task_id>")
def download_audio(task_id):
    task = tasks.get(task_id)
    if task and task["status"] == "completed":
        return send_file(task["file"])
    return "Not ready yet", 404