#!/opt/aime-venv/bin/python
"""CallRec persistent transcriber daemon.

Loads the faster-whisper large-v3-turbo CT2 model (decoder, CPU) ONCE and keeps
it warm for the whole callrec lifetime — this kills the ~22s CT2 load + ~14s
import cost that the old per-call `transcribe_cli.py` paid on every recording.

The NPU turbo ENCODER (rknn, 1.35 GB NPU memory) is acquired per request and
RELEASED right after, because the RK3588 NPU cannot hold the encoder and the
Qwen-3B LLM (rkllm_server) at the same time (proven: mali OOM). Freeing it here
lets rkllm load for the analysis that follows. The rknn file stays in the page
cache, so re-init after the first request is a few seconds, not the cold ~17s.

Protocol (line based, like spk_embed / rkllm_server):
  stdin : one absolute wav path per line
  stdout: "READY" once, then per request "OK <base64(json)>" or "ERR <base64(msg)>"
The JSON is base64-wrapped so newlines inside it never break the line protocol.
"""
import sys
import json
import base64
import time

import numpy as np

sys.path.insert(0, "/opt/callrec")

MODEL = "large-v3-turbo"
ROOT = "/opt/callrec/models/whisper"
RKNN = "/opt/callrec/models/turbo-encoder-fp16.rknn"

from faster_whisper import WhisperModel
from faster_whisper.transcribe import get_ctranslate2_storage
from inference.rknpu_whisper_encoder import RknnWhisperEncoder

# CT2 decoder loaded ONCE, stays warm.
_model = WhisperModel(MODEL, device="cpu", compute_type="int8",
                      download_root=ROOT, cpu_threads=4)


def transcribe(wav):
    t_load = time.time()
    npu = RknnWhisperEncoder(RKNN, lib_path="/usr/lib/librknnrt.so")
    load_s = time.time() - t_load

    def npu_encode(features):
        feats = np.asarray(features, dtype=np.float32)
        return get_ctranslate2_storage(np.ascontiguousarray(npu.run(feats)))

    _model.encode = npu_encode
    try:
        t0 = time.time()
        segments, info = _model.transcribe(
            wav, language="de", beam_size=1, word_timestamps=True, vad_filter=True)
        words = []
        texts = []
        for seg in segments:
            texts.append(seg.text.strip())
            for w in (seg.words or []):
                words.append({
                    "word": w.word.strip(),
                    "t_start_ms": int(round((w.start or 0) * 1000)),
                    "t_end_ms": int(round((w.end or 0) * 1000)),
                    "probability": round(w.probability, 3),
                })
        dur = float(info.duration or 0)
        trans_s = time.time() - t0
        return {
            "text": " ".join(texts).strip(),
            "words": words,
            "duration": round(dur, 1),
            "load_s": round(load_s, 1),
            "transcribe_s": round(trans_s, 1),
            "rtf": round(trans_s / max(dur, 0.1), 3),
            "compute": "npu",
        }
    finally:
        # Release the NPU encoder so rkllm can load for the analysis step.
        npu.close()


def main():
    print("READY", flush=True)
    for line in sys.stdin:
        wav = line.strip()
        if not wav:
            continue
        try:
            out = transcribe(wav)
            payload = base64.b64encode(
                json.dumps(out, ensure_ascii=False).encode("utf-8")).decode("ascii")
            print("OK " + payload, flush=True)
        except Exception as e:  # noqa: BLE001
            msg = base64.b64encode(str(e).encode("utf-8")).decode("ascii")
            print("ERR " + msg, flush=True)


if __name__ == "__main__":
    main()