"""
Optimised NeMo Parakeet-TDT streaming demo for CPU-only Hugging Face Spaces
"""

import os, time, threading, queue, logging, re
import numpy as np
import gradio as gr
from scipy import signal
import torch
from nemo.collections.asr.models import ASRModel

# ────────────────────────────────────────────────
# General CPU settings (2 vCPU space)
# ────────────────────────────────────────────────
os.environ["OMP_NUM_THREADS"] = "2"          # One MKL/OpenMP thread per vCPU
torch.set_num_threads(2)
torch.backends.quantized.engine = "fbgemm"   # Fastest INT8 kernels on x86

# ────────────────────────────────────────────────
# Logging
# ────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("asr_app")

# ────────────────────────────────────────────────
# Constants
# ────────────────────────────────────────────────
SR              = 16_000
CHUNK_SECONDS   = 4
CHUNK_SAMPLES   = SR * CHUNK_SECONDS

# ────────────────────────────────────────────────
# Prepare UI Description data from README.md file
# ────────────────────────────────────────────────
with open('README.md', 'r', encoding='utf-8') as file:
    content = file.read()
README_CONTENT_without_YAML = re.sub(r'^---.*?---\s*', '', content, flags=re.DOTALL)

# ────────────────────────────────────────────────
# ASR Application
# ────────────────────────────────────────────────
class ASRApp:
    def __init__(self):
        self.audio_queue      = queue.Queue(maxsize=8)
        self.transcript_queue = queue.Queue()
        self.transcript_list  = []
        self._load_model()
        self._start_worker()

    # ---------- helpers ----------
    def _log(self, func: str, msg: str):
        logger.info(
            f"{func} | audio_q={self.audio_queue.qsize():02}, "
            f"txt_q={self.transcript_queue.qsize():02} | {msg}"
        )

    # ---------- model ----------
    def _load_model(self):
        self._log("load_model", "loading Parakeet-TDT-0.6B-V2 (CPU)…")
        t0 = time.time()
        model = ASRModel.from_pretrained(
            model_name="nvidia/parakeet-tdt-0.6b-v2",
            map_location="cpu",
        )
        model.eval()
        self.asr_model = model
        self._log("load_model", f"model ready in {time.time()-t0:.1f}s")
        with torch.inference_mode():
            _ = self.asr_model.transcribe([np.zeros(SR, dtype=np.float32)])
        self._log("load_model", "warm-up done")
    
    # ---------- threading ----------
    def _start_worker(self):
        threading.Thread(target=self._worker, daemon=True).start()

    def _worker(self):
        buf = np.array([], dtype=np.float32)
        while True:
            try:
                while len(buf) < CHUNK_SAMPLES:
                    buf = np.concatenate([buf, self.audio_queue.get()])
                chunk, buf = buf[:CHUNK_SAMPLES], buf[CHUNK_SAMPLES:]
                self._log("_worker", f"→ transcribe {len(chunk)} samples")
                t0 = time.time()
                with torch.inference_mode():
                    out = self.asr_model.transcribe([chunk])
                dur = time.time() - t0
                text = out[0].text
                self._log("_worker", f"inference {dur:.2f}s → “{text}”")
                self.transcript_queue.put(text)
            except Exception as e:
                self._log("_worker", f"ASR error: {e}")

    # ---------- audio preprocessing ----------
    def _preprocess(self, audio):
        sr, y = audio
        y = signal.resample_poly(y, SR, sr)
        y = y.astype(np.float32)
        y /= (np.abs(y).max() + 1e-9)
        return y

    # ---------- Gradio stream callback ----------
    def stream_fn(self, audio):
        self._log("stream_fn", "audio arrived")
        self.audio_queue.put(self._preprocess(audio))
        while not self.transcript_queue.empty():
            self.transcript_list.append(self.transcript_queue.get())
        return (
            " ".join(self.transcript_list)
            if self.transcript_list
            else "…listening…"
        )

# ────────────────────────────────────────────────
# Gradio UI
# ────────────────────────────────────────────────
asr_app = ASRApp()
with gr.Blocks() as demo:
    mic = gr.Audio(
        sources=["microphone"],
        type="numpy",
        streaming=True,
        label="Microphone",
    )
    out = gr.Textbox(label="Transcription")
    gr.Markdown(README_CONTENT_without_YAML)
    
    mic.stream(
        fn=asr_app.stream_fn,
        inputs=mic,
        outputs=out,
        stream_every=0.5,
    )

asr_app._log("main", "launching UI")
demo.launch()