"""Small speech-to-text server run by the wxch agent inside the engine container (started by wxch-runner).

POST /transcribe?language=auto|th|en|...&task=transcribe|translate   body = the audio or video file as is
-> {"text": "...", "language": "...", "segments": [{"start": 0.0, "end": 2.5, "text": "..."}]}
GET /health -> 200 as soon as the server is up ("model_loaded" says whether the model is in memory yet).

One request at a time (the GPU is shared); the agent queues jobs.

The model is loaded when the first file arrives, not when the server starts: starting the engine (or warming it up for a page that
is open) must not take the memory yet, and the agent makes room in memory just before it sends the file.
"""
import argparse
import io
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlparse

import av
import numpy as np
import torch
from transformers import pipeline

MODEL = "/work/models/whisper-large-v3-turbo"
MAX_BYTES = 120 * 1024 * 1024
lock = threading.Lock()
pipe = None
model_path = MODEL


def get_pipe():
    """Load the speech model the first time it is needed (the caller holds the lock)."""
    global pipe
    if pipe is None:
        dev = "cuda:0" if torch.cuda.is_available() else "cpu"
        pipe = pipeline("automatic-speech-recognition", model=model_path, device=dev,
                        torch_dtype=torch.float16 if dev != "cpu" else torch.float32)
    return pipe


def decode(data):
    """Any audio or video container to 16 kHz mono float32."""
    out = []
    with av.open(io.BytesIO(data)) as c:
        stream = next((s for s in c.streams if s.type == "audio"), None)
        if stream is None:
            raise ValueError("the file has no audio")
        rs = av.audio.resampler.AudioResampler(format="flt", layout="mono", rate=16000)
        for frame in c.decode(stream):
            for f in rs.resample(frame):
                out.append(f.to_ndarray().reshape(-1))
        for f in rs.resample(None) or []:
            out.append(f.to_ndarray().reshape(-1))
    if not out:
        raise ValueError("the audio is empty")
    return np.concatenate(out).astype(np.float32)


def transcribe(data, language, task):
    audio = decode(data)
    kw = {"task": task}
    if language != "auto":
        kw["language"] = language
    with lock:
        res = get_pipe()({"raw": audio, "sampling_rate": 16000}, return_timestamps=True, chunk_length_s=30,
                         generate_kwargs=kw)
    segs = []
    for ch in res.get("chunks", []):
        a, b = ch["timestamp"]
        segs.append({"start": float(a or 0), "end": float(b if b is not None else (a or 0)), "text": ch["text"].strip()})
    return {"text": res["text"].strip(), "language": language, "segments": segs}


class H(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def send_json(self, code, obj):
        b = json.dumps(obj, ensure_ascii=False).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def do_GET(self):
        if urlparse(self.path).path == "/health":
            return self.send_json(200, {"ok": True, "model_loaded": pipe is not None})
        self.send_json(404, {"error": "not found"})

    def do_POST(self):
        u = urlparse(self.path)
        if u.path != "/transcribe":
            return self.send_json(404, {"error": "not found"})
        q = parse_qs(u.query)
        language = (q.get("language", ["auto"])[0] or "auto")[:8]
        task = q.get("task", ["transcribe"])[0]
        if task not in ("transcribe", "translate"):
            return self.send_json(400, {"error": "bad task"})
        n = int(self.headers.get("Content-Length", "0"))
        if n <= 0 or n > MAX_BYTES:
            return self.send_json(413, {"error": "file too large or empty"})
        data = self.rfile.read(n)
        try:
            self.send_json(200, transcribe(data, language, task))
        except Exception as e:  # report to the agent instead of a bare 500
            self.send_json(500, {"error": str(e)[:500]})


def main():
    global model_path
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8189)
    ap.add_argument("--host", default="0.0.0.0")  # inside a container; an engine run directly on the machine passes 127.0.0.1
    ap.add_argument("--model", default=MODEL)
    a = ap.parse_args()
    model_path = a.model
    ThreadingHTTPServer((a.host, a.port), H).serve_forever()


if __name__ == "__main__":
    main()