"""Duplex turn-taking: Silero VAD (MIT) per 32 ms frame, end-of-turn by trailing silence, barge-in while the agent
speaks/thinks, and light backchannel cues at mid-turn pauses. Pure frame clock (16 kHz, 512-sample frames)."""
import torch, numpy as np
FR = 512; SR = 16000; FMS = FR / SR

class TurnDetector:
    def __init__(self, eot_ms=600, start_ms=96, barge_ms=192, on=0.5, off=0.35, preroll_ms=300,
                 bc_min_turn_s=4.5, bc_pause_ms=(288, 448), bc_gap_s=10.0, spec_ms=256):
        from silero_vad import load_silero_vad
        self.vad = load_silero_vad()
        self.eot_f = int(eot_ms / 1000 / FMS); self.start_f = max(1, int(start_ms / 1000 / FMS)); self.barge_f = int(barge_ms / 1000 / FMS)
        self.on, self.off = on, off; self.pre_f = int(preroll_ms / 1000 / FMS)
        self.spec_f = int(spec_ms / 1000 / FMS) if spec_ms else 0
        self.bc_min, self.bc_pause, self.bc_gap = bc_min_turn_s, (int(bc_pause_ms[0] / 1000 / FMS), int(bc_pause_ms[1] / 1000 / FMS)), bc_gap_s
        self.reset()

    def reset(self):
        self.vad.reset_states(); self.t = 0.0; self.in_speech = False; self.on_count = 0; self.sil = 0
        self.ring = []; self.buf = []; self.speech_frames = 0; self.onset_t = None; self.last_bc = -99; self.bc_done_this_pause = False
        self.barged = False; self.busy_speech = 0; self._was_busy = False

    def feed(self, frame, agent_busy=False):
        """frame: float32[512]. returns list of (event, info). events: start, barge, backchannel, eot"""
        ev = []
        p = float(self.vad(torch.from_numpy(np.ascontiguousarray(frame, dtype=np.float32)), SR).item())
        self.t += FMS
        sp = p > (self.off if self.in_speech else self.on)
        if agent_busy and not self._was_busy:
            self.busy_speech = 0; self.barged = False
        self._was_busy = agent_busy
        if agent_busy:
            self.busy_speech = self.busy_speech + 1 if p > self.on else max(0, self.busy_speech - 1)
            if not self.barged and self.busy_speech >= self.barge_f:
                self.barged = True; ev.append(("barge", dict(t=self.t, onset=self.onset_t)))
        if not self.in_speech:
            self.ring.append(frame); self.ring = self.ring[-self.pre_f:]
            self.on_count = self.on_count + 1 if sp else 0
            if self.on_count >= self.start_f:
                self.in_speech = True; self.sil = 0; self.speech_frames = self.on_count
                self.onset_t = self.t - self.on_count * FMS; self.buf = list(self.ring); self.bc_done_this_pause = False
                ev.append(("start", dict(t=self.onset_t, p=p)))
            if agent_busy and self.on_count >= self.barge_f:
                pass
        else:
            self.buf.append(frame)
            if sp:
                if self.spec_f and self.sil >= self.spec_f:
                    ev.append(("resume", dict(t=self.t)))          # speculative response must be dropped
                self.sil = 0; self.speech_frames += 1; self.bc_done_this_pause = False
            else:
                self.sil += 1
                if self.spec_f and self.sil == self.spec_f:
                    keep = len(self.buf) - self.sil + int(0.2 / FMS)
                    ev.append(("spec", dict(t=self.t, speech_end=self.t - self.sil * FMS, audio=np.concatenate(self.buf[:max(1, keep)]))))
            dur = self.speech_frames * FMS
            if (not agent_busy and dur >= self.bc_min and self.bc_pause[0] <= self.sil <= self.bc_pause[1]
                    and not self.bc_done_this_pause and self.t - self.last_bc >= self.bc_gap):
                self.bc_done_this_pause = True; self.last_bc = self.t; ev.append(("backchannel", dict(t=self.t, turn_s=dur)))
            if self.sil >= self.eot_f:
                end_t = self.t - self.sil * FMS
                keep = len(self.buf) - self.sil + int(0.2 / FMS)
                audio = np.concatenate(self.buf[:max(1, keep)])
                ev.append(("eot", dict(t=self.t, speech_end=end_t, onset=self.onset_t, audio=audio, speech_s=dur)))
                self.in_speech = False; self.on_count = 0; self.buf = []; self.ring = []
        return ev

    def partial_audio(self):
        return np.concatenate(self.buf) if self.buf else np.zeros(0, np.float32)