"""Smart Turn Detection — Hugging Face Space demo (ONNX, CPU).
Runs the exported 8s model: raw 16k waveform -> P(turn complete).
Inference is CPU-only, so predictions don't touch the GPU queue. The one decorated
`_gpu_probe` below exists only so the ZeroGPU runtime accepts the Space; it's never called."""
import numpy as np
import onnxruntime as ort
import gradio as gr

try:
    import spaces

    @spaces.GPU
    def _gpu_probe():  # satisfies ZeroGPU's "must expose a GPU function" check; unused
        return None
except Exception:
    pass

SR, WIN = 16000, 8
N = SR * WIN  # 128000
SESS = ort.InferenceSession("smart_turn_8s.onnx", providers=["CPUExecutionProvider"])


def prob_complete(wav, sr):
    wav = np.asarray(wav, dtype=np.float32)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    if sr != SR:
        import librosa
        wav = librosa.resample(wav, orig_sr=sr, target_sr=SR)
    mx = np.abs(wav).max()
    if mx > 0:
        wav = wav / mx
    x = wav[-N:]                                  # last 8s
    buf = np.zeros((1, N), np.float32)
    buf[0, N - len(x):] = x                       # right-align (left-pad)
    return float(SESS.run(None, {"waveform": buf})[0][0])


def predict(audio, threshold):
    if audio is None:
        return "Record or upload some speech.", None
    sr, wav = audio
    if wav.dtype.kind in "iu":
        wav = wav.astype(np.float32) / np.iinfo(wav.dtype).max
    p = prob_complete(wav, sr)
    verdict = "✅ Turn COMPLETE — safe to respond" if p > threshold else "⏳ Still speaking — wait"
    return f"{verdict}\n\nP(complete) = {p:.3f}   (threshold {threshold:.2f})", {"complete": p, "still speaking": 1 - p}


with gr.Blocks(title="Smart Turn Detection") as demo:
    gr.Markdown(
        "# 🎙️ Smart Turn Detection (whisper-tiny, 8s)\n"
        "Has the speaker **finished their turn**, or are they just pausing? Audio only, "
        "no transcription. Tuned for English + Hindi/Hinglish. ~7.8M params, ~38 ms/clip on CPU.\n\n"
        "_Try a finished sentence vs. trailing off with 'umm…'. Raise the threshold to interrupt less._")
    inp = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Speech")
    thr = gr.Slider(0.05, 0.95, value=0.5, step=0.05, label="Decision threshold")
    btn = gr.Button("Detect turn", variant="primary")
    out_txt = gr.Textbox(label="Verdict")
    out_lbl = gr.Label(label="Probabilities")
    btn.click(predict, [inp, thr], [out_txt, out_lbl])
    inp.stop_recording(predict, [inp, thr], [out_txt, out_lbl])

if __name__ == "__main__":
    demo.launch()