#!/usr/bin/env python3
"""CPU-only Piper service shared with the routing Jetson."""

from __future__ import annotations

import io
import json
import os
import threading
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from piper import PiperVoice, SynthesisConfig


HOST = os.getenv("HERMES_TTS_HOST", "0.0.0.0")
PORT = int(os.getenv("HERMES_TTS_PORT", "9001"))
VOICE_NAME = os.getenv("HERMES_TTS_VOICE", "en_US-lessac-high")
CACHE_DIR = Path(os.getenv("HERMES_TTS_CACHE", "/cache/piper"))
MAX_TEXT_CHARS = 5000
VOICE_LOCK = threading.Lock()


def _json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.send_header("Cache-Control", "no-store")
    handler.end_headers()
    handler.wfile.write(body)


class SpeechHandler(BaseHTTPRequestHandler):
    """Serve health and bounded local speech synthesis."""

    server_version = "AtlasPiper/1"

    def log_message(self, message: str, *args: object) -> None:
        print(f"[tts] {self.address_string()} {message % args}", flush=True)

    def do_GET(self) -> None:
        if self.path != "/health":
            _json(self, 404, {"error": "not found"})
            return
        _json(self, 200, {"ok": True, "voice": VOICE_NAME, "device": "cpu"})

    def do_POST(self) -> None:
        if self.path != "/v1/audio/speech":
            _json(self, 404, {"error": "not found"})
            return
        content_length = int(self.headers.get("Content-Length", "0") or "0")
        if content_length <= 0 or content_length > 64 * 1024:
            _json(self, 413, {"error": "request is missing or too large"})
            return
        try:
            payload = json.loads(self.rfile.read(content_length).decode("utf-8"))
            text = str(payload.get("input") or payload.get("text") or "").strip()
            speed = float(payload.get("speed") or 1.0)
        except (TypeError, ValueError, json.JSONDecodeError):
            _json(self, 400, {"error": "invalid JSON request"})
            return
        if not text:
            _json(self, 400, {"error": "input is required"})
            return
        if len(text) > MAX_TEXT_CHARS:
            _json(self, 400, {"error": "input is too long"})
            return
        speed = min(2.0, max(0.5, speed))

        output = io.BytesIO()
        try:
            with VOICE_LOCK, wave.open(output, "wb") as wav_file:
                self.server.voice.synthesize_wav(  # type: ignore[attr-defined]
                    text,
                    wav_file,
                    SynthesisConfig(length_scale=1.0 / speed),
                )
            audio = output.getvalue()
            self.send_response(200)
            self.send_header("Content-Type", "audio/wav")
            self.send_header("Content-Length", str(len(audio)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(audio)
        except Exception as exc:
            print(f"[tts] synthesis failed: {exc}", flush=True)
            _json(self, 500, {"error": "speech synthesis failed"})


def main() -> None:
    """Load the checksum-pinned voice from the image and serve it on CPU."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    model_path = CACHE_DIR / f"{VOICE_NAME}.onnx"
    config_path = CACHE_DIR / f"{VOICE_NAME}.onnx.json"
    if not model_path.exists() or not config_path.exists():
        raise RuntimeError(f"baked Piper voice is missing: {VOICE_NAME}")
    print(f"[tts] loading Piper voice {VOICE_NAME} on CPU", flush=True)
    voice = PiperVoice.load(model_path, config_path, use_cuda=False, download_dir=CACHE_DIR)
    server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
    server.voice = voice  # type: ignore[attr-defined]
    print(f"[tts] ready on {HOST}:{PORT}", flush=True)
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()