"""Qwen3-TTS HTTP server — OpenAI-compatible /v1/audio/speech endpoint."""
import argparse, io, os, sys
from flask import Flask, request, send_file, jsonify

app = Flask(__name__)
MODEL_DIR = "/models/qwen3-tts"
engine = None

SPEAKERS = ["Vivian","Serena","uncle_fu","Dylan","Eric","Ryan","Aiden","ono_anna","Sohee"]
LANGUAGES = ["english","chinese","japanese","korean","german","spanish","french","russian","italian","portuguese"]

def init_engine():
    global engine
    # TTSEngine uses Path.cwd() as project_root and builds relative paths from it.
    # model_dir must be a subdirectory of cwd for relative_to() to work.
    parent = os.path.dirname(os.path.abspath(MODEL_DIR))
    basename = os.path.basename(MODEL_DIR)
    os.chdir(parent)
    from qwen3_tts_gguf.inference import TTSEngine
    engine = TTSEngine(model_dir=basename, onnx_provider="CPUExecutionProvider")
    if not engine.ready:
        print(f"[server] WARNING: engine created but not ready", flush=True)

@app.route("/health")
def health():
    ready = engine is not None and engine.ready
    return jsonify({"status": "ok" if ready else "loading", "speakers": SPEAKERS, "languages": LANGUAGES})

@app.route("/v1/audio/speech", methods=["POST"])
def speech():
    if engine is None or not engine.ready:
        return jsonify({"error": "engine not ready"}), 503
    data = request.get_json(force=True)
    text = data.get("input", "")
    if not text:
        return jsonify({"error": "missing input"}), 400
    speaker = data.get("voice", "Vivian")
    if speaker not in SPEAKERS:
        speaker = "Vivian"
    language = data.get("language", "english")
    if language not in LANGUAGES:
        language = "english"
    instruct = data.get("instruct", "")
    try:
        import numpy as np, wave as _wave
        from qwen3_tts_gguf.inference import TTSConfig
        stream = engine.create_stream(n_ctx=2048)
        if stream is None:
            return jsonify({"error": "failed to create stream"}), 500
        cfg = TTSConfig()
        result = stream.custom(text=text, speaker=speaker, language=language,
                               instruct=instruct, config=cfg)
        stream.join()
        if result is None:
            return jsonify({"error": "synthesis returned None"}), 500
        audio = result.audio if hasattr(result, 'audio') and result.audio is not None else None
        if audio is None:
            tmp = os.path.join("/tmp", "tts_out.wav")
            result.save(tmp)
            import soundfile as sf
            audio, _ = sf.read(tmp, dtype='float32')
            os.remove(tmp)
        buf = io.BytesIO()
        with _wave.open(buf, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(24000)
            pcm = (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16).tobytes()
            wf.writeframes(pcm)
        buf.seek(0)
        return send_file(buf, mimetype="audio/wav", download_name="speech.wav")
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8072)
    p.add_argument("--model-dir", default=MODEL_DIR)
    a = p.parse_args()
    MODEL_DIR = a.model_dir
    init_engine()
    app.run(host=a.host, port=a.port)