#!/opt/aime-venv/bin/python
"""CallRec NPU transcriber: wav -> JSON words on stdout.
large-v3-turbo decoder on CPU (int8), encoder on RK3588 NPU (rknn fp16).
Output matches what recorder.js consumes: {"text","words":[{word,t_start_ms,t_end_ms,probability}]}.
Usage: transcribe_cli.py <wav> [--cpu]   (--cpu = no NPU, plain CPU encoder for A/B)
"""
import sys, json, time
import numpy as np

WAV = sys.argv[1]
USE_NPU = "--cpu" not in sys.argv[2:]
import os
BATCH = int(os.environ.get("WHISPER_BATCH", "8"))
MODEL = "large-v3-turbo"
ROOT = "/opt/callrec/models/whisper"
RKNN = "/opt/callrec/models/turbo-encoder-fp16.rknn"
sys.path.insert(0, "/opt/callrec")

from faster_whisper import WhisperModel, BatchedInferencePipeline

t0 = time.time()
model = WhisperModel(MODEL, device="cpu", compute_type="int8",
                     download_root=ROOT, cpu_threads=4)

if USE_NPU:
    from faster_whisper.transcribe import get_ctranslate2_storage
    from inference.rknpu_whisper_encoder import RknnWhisperEncoder
    npu = RknnWhisperEncoder(RKNN, lib_path="/usr/lib/librknnrt.so")
    def npu_encode(features):
        feats = np.asarray(features, dtype=np.float32)
        out = npu.run(feats)
        return get_ctranslate2_storage(np.ascontiguousarray(out))
    model.encode = npu_encode

engine = BatchedInferencePipeline(model=model) if BATCH > 1 else model
t_load = time.time() - t0
t0 = time.time()
if BATCH > 1:
    segments, info = engine.transcribe(WAV, language="de", beam_size=1,
                                       batch_size=BATCH, word_timestamps=True)
else:
    segments, info = engine.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),
        })
t_trans = time.time() - t0
out = {
    "text": " ".join(texts).strip(),
    "words": words,
    "duration": round(float(info.duration or 0), 1),
    "load_s": round(t_load, 1),
    "transcribe_s": round(t_trans, 1),
    "rtf": round(t_trans / max(float(info.duration or 0.1), 0.1), 3),
    "compute": "npu" if USE_NPU else "cpu",
    "batch": BATCH,
}
print(json.dumps(out, ensure_ascii=False))