"""
Streaming ASR / VAD endpointing tests.

Uses synthetic audio and a mock transcriber so it runs with no models and
no network. Verifies the state machine, not Whisper.
"""

import sys
import logging
import numpy as np

logging.basicConfig(level=logging.WARNING)

from vad import load_vad, FRAME_MS, SAMPLE_RATE
from streaming_asr import StreamingASR, EndpointConfig, State, _is_hallucination

PASS, FAIL = "\033[92mPASS\033[0m", "\033[91mFAIL\033[0m"
results = []


def check(name, cond, detail=""):
    results.append((name, cond))
    print(f"  [{PASS if cond else FAIL}] {name}" + (f"  — {detail}" if detail else ""))


# ── Synthetic audio ───────────────────────────────────────────────────────────

def silence(ms, noise=0.0005):
    n = int(SAMPLE_RATE * ms / 1000)
    return (np.random.randn(n) * noise).astype(np.float32)


def speech(ms, amp=0.25):
    """Voice-like: 120Hz glottal pulse + formants + amplitude modulation."""
    n = int(SAMPLE_RATE * ms / 1000)
    t = np.arange(n) / SAMPLE_RATE
    sig = (np.sin(2 * np.pi * 120 * t)
           + 0.5 * np.sin(2 * np.pi * 700 * t)
           + 0.3 * np.sin(2 * np.pi * 1220 * t))
    envelope = 0.6 + 0.4 * np.sin(2 * np.pi * 4 * t)      # syllable rate
    sig = sig * envelope
    sig += np.random.randn(n) * 0.01
    return (sig / np.abs(sig).max() * amp).astype(np.float32)


def feed(sasr, audio, chunk_ms=100):
    """Feed audio in realistic small chunks, collecting all events."""
    events = []
    step = int(SAMPLE_RATE * chunk_ms / 1000)
    for i in range(0, len(audio), step):
        events.extend(sasr.accept_audio(audio[i:i + step], SAMPLE_RATE))
    return events


# Mock transcriber records exactly what audio it was handed
DECODE_LOG = []


def mock_transcribe(audio, sr):
    DECODE_LOG.append(len(audio) / sr * 1000)      # duration in ms
    return f"transcript_of_{len(audio)/sr:.2f}s"


def new_asr(**kw):
    DECODE_LOG.clear()
    cfg = EndpointConfig(**kw) if kw else EndpointConfig()
    return StreamingASR(transcribe_fn=mock_transcribe, config=cfg,
                        vad_backend="energy")


# ══════════════════════════════════════════════════════════════════════════════

print("\n══ VAD discriminates speech from silence ══")
vad = load_vad("energy")
for _ in range(12):                                   # let the floor settle
    vad.speech_prob(silence(FRAME_MS)[:512])
sil_probs = [vad.speech_prob(silence(FRAME_MS)[:512]) for _ in range(20)]
sp = speech(1000)
sp_probs = [vad.speech_prob(sp[i:i+512]) for i in range(0, 512*20, 512)]
check("silence scores low", np.mean(sil_probs) < 0.3,
      f"mean={np.mean(sil_probs):.2f}")
check("speech scores high", np.mean(sp_probs) > 0.6,
      f"mean={np.mean(sp_probs):.2f}")

print("\n══ Basic endpointing: speech → silence → ONE final ══")
sasr = new_asr()
audio = np.concatenate([silence(500), speech(1200), silence(1200)])
events = feed(sasr, audio)
finals = [e for e in events if e.kind == "final"]
check("exactly one final emitted", len(finals) == 1, f"got {len(finals)}")
check("speech_start fired", any(e.kind == "speech_start" for e in events))
check("returns to IDLE after endpoint", sasr.state is State.IDLE)

print("\n══ Preroll: first phoneme is not clipped ══")
sasr = new_asr(preroll_ms=300)
audio = np.concatenate([silence(400), speech(1000), silence(1200)])
feed(sasr, audio)
check("decoded audio longer than speech alone (preroll included)",
      DECODE_LOG and DECODE_LOG[-1] > 1000,
      f"decoded {DECODE_LOG[-1]:.0f}ms for 1000ms of speech")

print("\n══ Hangover: a natural pause does NOT end the turn ══")
sasr = new_asr(endpoint_silence_ms=700)
# 400ms pause is shorter than the 700ms endpoint → must stay ONE utterance
audio = np.concatenate([silence(400), speech(700), silence(400),
                        speech(700), silence(1200)])
events = feed(sasr, audio)
finals = [e for e in events if e.kind == "final"]
check("short internal pause does not split the turn", len(finals) == 1,
      f"got {len(finals)} finals")

print("\n══ Long pause DOES split into two turns ══")
sasr = new_asr(endpoint_silence_ms=700)
audio = np.concatenate([silence(400), speech(700), silence(1300),
                        speech(700), silence(1300)])
events = feed(sasr, audio)
finals = [e for e in events if e.kind == "final"]
check("long pause produces two separate turns", len(finals) == 2,
      f"got {len(finals)} finals")

print("\n══ Silence is never sent to Whisper (hallucination guard) ══")
sasr = new_asr()
feed(sasr, silence(4000))
check("no decode call on pure silence", len(DECODE_LOG) == 0,
      f"{len(DECODE_LOG)} decode calls")
check("no final event on pure silence",
      not any(e.kind == "final" for e in feed(sasr, silence(2000))))

print("\n══ Cough / click is discarded, not transcribed ══")
sasr = new_asr(min_speech_ms=250)
audio = np.concatenate([silence(400), speech(90), silence(1500)])
events = feed(sasr, audio)
check("sub-threshold blip produces no final",
      not any(e.kind == "final" for e in events))
check("no decode call for the blip", len(DECODE_LOG) == 0,
      f"{len(DECODE_LOG)} calls")

print("\n══ Live partials during a long utterance ══")
sasr = new_asr(partial_interval_ms=600)
audio = np.concatenate([silence(400), speech(3500), silence(1200)])
events = feed(sasr, audio)
partials = [e for e in events if e.kind == "partial"]
check("partials emitted while speaking", len(partials) >= 2,
      f"got {len(partials)} partials")
check("final still emitted after partials",
      any(e.kind == "final" for e in events))
check("partials precede the final",
      events.index(next(e for e in events if e.kind == "final"))
      > events.index(partials[0]) if partials else False)

print("\n══ Barge-in: caller interrupts the agent ══")
sasr = new_asr(bargein_speech_ms=220)
sasr.agent_speaking = True
events = feed(sasr, np.concatenate([silence(300), speech(900)]))
check("barge-in event raised", any(e.kind == "bargein" for e in events))
check("agent_speaking cleared", sasr.agent_speaking is False)

print("\n══ Barge-in does NOT fire on background noise ══")
sasr = new_asr(bargein_speech_ms=220)
sasr.agent_speaking = True
events = feed(sasr, silence(2000, noise=0.002))
check("no barge-in on quiet background",
      not any(e.kind == "bargein" for e in events))

print("\n══ Max-duration force endpoint ══")
sasr = new_asr(max_utterance_ms=2000)
events = feed(sasr, np.concatenate([silence(300), speech(6000)]))
check("non-stop speaker is force-endpointed",
      any(e.kind == "final" for e in events))

print("\n══ flush() finalizes a turn in progress (caller hung up) ══")
sasr = new_asr()
feed(sasr, np.concatenate([silence(400), speech(1000)]))
events = sasr.flush()
check("flush emits the pending final",
      any(e.kind == "final" for e in events))

print("\n══ Chunk-size independence (frame alignment) ══")
for chunk_ms in (20, 33, 100, 250):
    sasr = new_asr()
    audio = np.concatenate([silence(400), speech(1200), silence(1200)])
    events = feed(sasr, audio, chunk_ms=chunk_ms)
    n = len([e for e in events if e.kind == "final"])
    check(f"chunk={chunk_ms}ms → 1 final", n == 1, f"got {n}")

print("\n══ Whisper hallucination filter ══")
check("'Thank you.' filtered", _is_hallucination("Thank you."))
check("'Subtitles by the Amara.org community' filtered",
      _is_hallucination("Subtitles by the Amara.org community"))
check("'♪' filtered", _is_hallucination("♪"))
check("real Hausa text NOT filtered",
      not _is_hallucination("Ina son duba asusuna"))

print("\n══ int16 input auto-converted ══")
sasr = new_asr()
audio_f = np.concatenate([silence(400), speech(1200), silence(1200)])
audio_i16 = (audio_f * 32767).astype(np.int16)
events = feed(sasr, audio_i16)
check("int16 stream produces a final",
      any(e.kind == "final" for e in events))

print("\n══ Resampling from 8kHz (telephony) ══")
sasr = new_asr()
audio = np.concatenate([silence(400), speech(1200), silence(1200)])
audio_8k = audio[::2]                                  # crude 8kHz
events = []
step = 800
for i in range(0, len(audio_8k), step):
    events.extend(sasr.accept_audio(audio_8k[i:i + step], 8000))
check("8kHz telephony audio endpoints correctly",
      any(e.kind == "final" for e in events))

# ── Summary ───────────────────────────────────────────────────────────────────
passed = sum(1 for _, ok in results if ok)
total = len(results)
print(f"\n{'='*62}\n  {passed}/{total} checks passed\n{'='*62}")
sys.exit(0 if passed == total else 1)