#!/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

import onnxruntime
from piper import PiperConfig, 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
ONNX_THREADS = max(1, int(os.getenv("HERMES_TTS_ONNX_THREADS", "4")))
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}")
    with config_path.open("r", encoding="utf-8") as config_file:
        config = PiperConfig.from_dict(json.load(config_file))
    session_options = onnxruntime.SessionOptions()
    session_options.intra_op_num_threads = ONNX_THREADS
    session_options.inter_op_num_threads = 1
    session = onnxruntime.InferenceSession(
        str(model_path),
        sess_options=session_options,
        providers=["CPUExecutionProvider"],
    )
    voice = PiperVoice(session=session, config=config, download_dir=CACHE_DIR)
    print(
        f"[tts] loaded Piper voice {VOICE_NAME} on CPU with {ONNX_THREADS} ONNX threads",
        flush=True,
    )
    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()