"""Minimal inference sidecar for tts-gateway (Phase 3, TTS_GATEWAY_PLAN.md).

Runs only SynthesizerTrn.infer() - no Flask, no training code, no Japanese
text handling (that's done entirely in Go by
internal/adapters/secondary/jtalk; this process receives already-normalized,
already-interspersed symbol IDs and just runs the model). Vendored model
code (models/, commons.py) is copied from tmp/reference/uma-tts-api,
including the transforms.py edits made during Phase 0's ONNX spike - those
are behavior-preserving no-ops for this eager-mode infer() path too (see
tmp/reference/uma-tts-api/spike/FINDINGS.md), so there was no reason to
maintain two divergent copies.
"""

import json
import logging
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import torch

from models import utils
from models.models import SynthesizerTrn

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

CONFIG_PATH = os.environ.get("CONFIG_PATH", "/models/uma.json")
CHECKPOINT_PATH = os.environ.get("CHECKPOINT_PATH", "/models/G_790000.pth")
PORT = int(os.environ.get("PORT", "50054"))

# len(text.symbols.symbols) from the reference repo: 1 (pad) + 8 (special) +
# 22 (punctuation) + 62 (letters+digits) + 84 (dummy) = 177. Hardcoded rather
# than importing text/symbols.py, since this sidecar deliberately carries no
# text-handling code - Go's internal/core/domain/symbols.go is the source of
# truth for the vocabulary now, and the two are verified to agree in
# tts-gateway's tests.
N_VOCAB = 177

model_lock = threading.Lock()


def load_model():
    hps = utils.get_hparams_from_file(CONFIG_PATH)
    net_g = SynthesizerTrn(
        n_vocab=N_VOCAB,
        spec_channels=hps.data.filter_length // 2 + 1,
        segment_size=hps.train.segment_size // hps.data.hop_length,
        n_speakers=hps.data.n_speakers,
        **hps.model,
    )
    net_g.eval()
    if torch.cuda.is_available():
        net_g.cuda()
    else:
        logger.warning("CUDA not available, running on CPU")
    utils.load_checkpoint(CHECKPOINT_PATH, net_g, None)
    logger.info("model loaded from %s", CHECKPOINT_PATH)
    return net_g, hps


net_g, hps = load_model()


def synthesize(symbol_ids, speaker_id, noise_scale, noise_scale_w, length_scale):
    x = torch.LongTensor(symbol_ids).unsqueeze(0)
    x_lengths = torch.LongTensor([len(symbol_ids)])
    sid = torch.LongTensor([speaker_id])
    if torch.cuda.is_available():
        x, x_lengths, sid = x.cuda(), x_lengths.cuda(), sid.cuda()

    with torch.no_grad(), model_lock:
        audio, *_ = net_g.infer(
            x, x_lengths, sid=sid,
            noise_scale=noise_scale, noise_scale_w=noise_scale_w, length_scale=length_scale,
        )
    return audio[0, 0].data.cpu().float().numpy()


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        logger.info("%s - %s", self.address_string(), fmt % args)

    def do_GET(self):
        if self.path == "/health":
            self._respond(200, b'{"status":"ok"}', "application/json")
        else:
            self._respond(404, b"not found", "text/plain")

    def do_POST(self):
        if self.path != "/synthesize":
            self._respond(404, b"not found", "text/plain")
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            body = json.loads(self.rfile.read(length))
            pcm = synthesize(
                body["symbol_ids"],
                body["speaker_id"],
                body.get("noise_scale", 0.37),
                body.get("noise_scale_w", 0.46),
                body.get("length_scale", 1.3),
            )
        except Exception as e:  # noqa: BLE001 - report any failure to the client
            logger.exception("synthesize failed")
            self._respond(400, str(e).encode("utf-8"), "text/plain")
            return

        payload = pcm.astype("<f4").tobytes()
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("X-Sample-Rate", str(hps.data.sampling_rate))
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _respond(self, code, body, content_type):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


if __name__ == "__main__":
    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
    logger.info("inference sidecar listening on :%d", PORT)
    server.serve_forever()