#!/usr/bin/env python3
"""OpenAI-compatible TTS API server for Kokoro on AX650 NPU."""
import argparse, json, tempfile, time
from http.server import HTTPServer, BaseHTTPRequestHandler
import soundfile as sf
from kokoro_axera import KokoroEngine

engine = None

class TTSHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        if self.path == "/v1/audio/speech":
            content_length = int(self.headers.get("Content-Length", 0))
            body = json.loads(self.rfile.read(content_length))
            text = body.get("input", "")
            t0 = time.time()
            audio, sr = engine.synthesize(text)
            elapsed = time.time() - t0

            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                sf.write(f.name, audio, sr)
                f.seek(0)
                data = f.read()

            self.send_response(200)
            self.send_header("Content-Type", "audio/wav")
            self.send_header("Content-Length", len(data))
            self.end_headers()
            self.wfile.write(data)
        elif self.path == "/health":
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b'{"status":"ok"}')
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        pass  # suppress logs

def main():
    global engine
    parser = argparse.ArgumentParser(description="Kokoro TTS API Server")
    parser.add_argument("--port", "-p", type=int, default=8880)
    parser.add_argument("--model-dir", type=str, default=None)
    args = parser.parse_args()

    print("Loading models...")
    engine = KokoroEngine(model_dir=args.model_dir)

    server = HTTPServer(("0.0.0.0", args.port), TTSHandler)
    print(f"Server running on http://0.0.0.0:{args.port}")
    print(f"POST /v1/audio/speech  {{'input': '<phonemes>', 'model': 'kokoro'}}")
    server.serve_forever()

if __name__ == "__main__":
    main()