"""
VSIP voice agent — a complete, runnable example.

    microphone ──► VSIP (turns, identity, barge-in)
                     │ turn verified (should_respond)
                     ▼
               OpenAI STT ──► OpenAI LLM (streamed) ──► OpenAI TTS (per sentence) ──► speakers
                     ▲                                                              │
                     └──────────── barge_in: stop playback + cancel ◄───────────────┘

Only the enrolled speaker is answered; anyone else is ignored. The enrolled
speaker can interrupt the agent mid-sentence; coughs and bystanders cannot.

Run with a microphone (use HEADPHONES — speaker echo confuses any voice agent):
    pip install -r requirements.txt
    export VSIP_API_KEY=vsip_...  OPENAI_API_KEY=sk-...
    python agent.py

Run without audio hardware (CI / smoke test) — stream WAV files instead:
    python agent.py --enroll-file me_intro.wav --input-file me_question.wav \
                    --no-playback --save-replies replies/ --duration 40
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import re
import sys
import time
import wave
from pathlib import Path
from typing import AsyncIterator, Optional

import numpy as np
from openai import AsyncOpenAI

from vsip import Turn, VSIPClient, VoiceSession, int16_to_wav

SAMPLE_RATE = 16000                 # VSIP wire format: 16 kHz mono int16
FRAME_SAMPLES = 320                 # 20 ms frames
TTS_SAMPLE_RATE = 24000             # OpenAI `pcm` speech output
ENROLL_PROMPT = (
    "Hello, my name is ___ and I'm setting up this assistant so it recognises my voice. "
    "I usually talk to it from my desk, sometimes with other people nearby. "
    "I'd like short, clear answers, and I'll interrupt when I need to."
)

# Models are configurable: OpenAI retires models on a schedule. As of Oct 2026,
# gpt-4o-mini-tts is removed 2027-01-06 and gpt-4o-mini on 2027-01-14 — set
# TTS_MODEL / LLM_MODEL to their successors before then.
STT_MODEL = os.environ.get("STT_MODEL", "gpt-transcribe")
LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")
TTS_MODEL = os.environ.get("TTS_MODEL", "gpt-4o-mini-tts")
TTS_VOICE = os.environ.get("TTS_VOICE", "alloy")
SYSTEM_PROMPT = os.environ.get(
    "SYSTEM_PROMPT",
    "You are a concise, friendly voice assistant. Answer in one to three short "
    "spoken sentences. If the user interrupted you, address what they just said.",
)

log = logging.getLogger("agent")


# ── Audio helpers ─────────────────────────────────────────────────────────────

def read_wav_16k(path: Path) -> np.ndarray:
    """Any PCM WAV → 16 kHz mono int16."""
    with wave.open(str(path), "rb") as w:
        rate, channels, width = w.getframerate(), w.getnchannels(), w.getsampwidth()
        raw = w.readframes(w.getnframes())
    if width != 2:
        raise SystemExit(f"{path}: only 16-bit PCM WAV is supported")
    audio = np.frombuffer(raw, dtype="<i2").astype(np.float32)
    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)
    if rate != SAMPLE_RATE:
        positions = np.arange(0, len(audio), rate / SAMPLE_RATE)
        audio = np.interp(positions, np.arange(len(audio)), audio)
    return np.clip(audio, -32768, 32767).astype("<i2")


def int16_to_float32_bytes(pcm: bytes) -> bytes:
    return (np.frombuffer(pcm, dtype="<i2").astype(np.float32) / 32768.0).tobytes()


class Speaker:
    """Cancellable playback of 24 kHz int16 PCM (or a timed no-op when muted)."""

    def __init__(self, enabled: bool, save_dir: Optional[Path]):
        self.enabled = enabled
        self.save_dir = save_dir
        self._cancel = asyncio.Event()
        self._stream = None
        self._saved = 0
        if save_dir:
            save_dir.mkdir(parents=True, exist_ok=True)
        if enabled:
            import sounddevice as sd

            self._stream = sd.RawOutputStream(samplerate=TTS_SAMPLE_RATE, channels=1, dtype="int16")
            self._stream.start()

    def cancel(self) -> None:
        self._cancel.set()

    def reset(self) -> None:
        self._cancel.clear()

    @property
    def cancelled(self) -> bool:
        return self._cancel.is_set()

    async def play(self, pcm: bytes) -> bool:
        """Play in 100 ms slices so a barge-in stops within ~100 ms. Returns
        False when cancelled."""
        if self.save_dir:
            self._saved += 1
            with wave.open(str(self.save_dir / f"reply_{self._saved:03d}.wav"), "wb") as w:
                w.setnchannels(1)
                w.setsampwidth(2)
                w.setframerate(TTS_SAMPLE_RATE)
                w.writeframes(pcm)
        step = TTS_SAMPLE_RATE // 10 * 2
        for offset in range(0, len(pcm), step):
            if self._cancel.is_set():
                return False
            chunk = pcm[offset:offset + step]
            if self._stream is not None:
                await asyncio.to_thread(self._stream.write, chunk)
            else:
                await asyncio.sleep(len(chunk) / 2 / TTS_SAMPLE_RATE)  # keep real-time pacing
        return True


# ── The agent ────────────────────────────────────────────────────────────────

class Agent:
    def __init__(self, client: VSIPClient, openai: AsyncOpenAI, speaker: Speaker, language: Optional[str]):
        self.client = client
        self.openai = openai
        self.speaker = speaker
        self.language = language
        self.session: Optional[VoiceSession] = None
        self.history: list[dict] = []
        self.locked = False

    # VSIP callbacks ----------------------------------------------------------

    async def on_turn(self, turn: Turn) -> None:
        seconds = len(turn.audio) / 2 / turn.sample_rate
        if not turn.should_respond:
            log.info("✗ %.1fs turn from %s — not the enrolled speaker, staying silent "
                     "(verified=%s, similarity=%s)", seconds, turn.speaker_id, turn.verified, turn.similarity)
            return
        started = time.perf_counter()
        text = await self.transcribe(turn.audio)
        if not text:
            return
        log.info("you%s: %s  (stt %.0f ms)", " [interrupted agent]" if turn.interrupted_agent else "",
                 text, (time.perf_counter() - started) * 1000)
        await self.respond(text, interrupted=turn.interrupted_agent, started=started)

    def on_barge_in(self, event) -> None:
        # VoiceSession already sent set_agent_state=idle; stop audio NOW.
        log.info("⚡ barge-in by %s — stopping playback", getattr(event, "speaker_id", "?"))
        self.speaker.cancel()

    def on_event(self, event) -> None:
        if event.type == "enrollment_result":
            raw = event.raw
            if raw.get("success"):
                log.info("enrolled '%s' · quality=%s · %s windows kept, %s rejected",
                         raw.get("label"), raw.get("quality", "n/a"),
                         raw.get("segments_used", "?"), raw.get("segments_rejected", "?"))
                if raw.get("quality") == "poor":
                    log.warning("voiceprint quality is poor — re-run enrollment somewhere quieter")
                if raw.get("locked") is not True and not self.locked:
                    # Older servers: lock in a separate step.
                    asyncio.get_running_loop().create_task(self.client.toggle_lock())
            else:
                log.error("enrollment failed: %s", raw.get("reason"))
        elif event.type == "lock_status":
            self.locked = str(event.raw.get("status", "")).startswith("locked:")
            log.info("lock: %s", event.raw.get("status"))
        elif event.type == "speaker_verification":
            log.info("verified %s → should_respond=%s (similarity=%s)", event.raw.get("speaker_id"),
                     event.raw.get("should_respond"), event.raw.get("similarity"))
        elif event.type == "error":
            log.error("VSIP error: %s", event.raw.get("message"))

    # STT / LLM / TTS ---------------------------------------------------------

    async def transcribe(self, pcm: bytes) -> str:
        kwargs = {}
        if self.language and STT_MODEL != "gpt-transcribe":
            kwargs["language"] = self.language
        result = await self.openai.audio.transcriptions.create(
            model=STT_MODEL, file=("turn.wav", int16_to_wav(pcm), "audio/wav"), **kwargs,
        )
        return (result.text or "").strip()

    async def sentences(self, messages: list[dict]) -> AsyncIterator[str]:
        stream = await self.openai.chat.completions.create(
            model=LLM_MODEL, messages=messages, stream=True, max_tokens=200,
        )
        buffer = ""
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                buffer += chunk.choices[0].delta.content
            parts = re.split(r"(?<=[.!?])\s+", buffer)
            for sentence in parts[:-1]:
                if sentence.strip():
                    yield sentence.strip()
            buffer = parts[-1]
        if buffer.strip():
            yield buffer.strip()

    async def speech(self, text: str) -> bytes:
        response = await self.openai.audio.speech.create(
            model=TTS_MODEL, voice=TTS_VOICE, input=text, response_format="pcm",
        )
        return response.content

    async def respond(self, text: str, *, interrupted: bool, started: float) -> None:
        user = {"role": "user", "content": f"[interrupting you] {text}" if interrupted else text}
        messages = [{"role": "system", "content": SYSTEM_PROMPT}, *self.history[-12:], user]
        audio_queue: asyncio.Queue = asyncio.Queue()
        spoken: list[str] = []
        self.speaker.reset()

        async def produce() -> None:
            # Synthesis of sentence N+1 overlaps playback of sentence N.
            try:
                async for sentence in self.sentences(messages):
                    if self.speaker.cancelled:
                        break
                    await audio_queue.put((sentence, await self.speech(sentence)))
            finally:
                await audio_queue.put(None)

        producer = asyncio.create_task(produce())
        speaking = False
        try:
            while True:
                item = await audio_queue.get()
                if item is None:
                    break
                sentence, pcm = item
                if not speaking:
                    assert self.session is not None
                    ready = time.perf_counter()
                    await self.session.speak_gate()     # never talk over an open user turn
                    await self.client.agent_speaking()   # arms VSIP barge-in detection
                    speaking = True
                    waited = time.perf_counter() - ready
                    log.info("first audio ready %.0f ms after the turn was verified%s",
                             (ready - started) * 1000,
                             f" (then waited {waited:.1f}s for someone to finish talking)" if waited > 0.2 else "")
                log.info("agent: %s", sentence)
                if not await self.speaker.play(pcm):
                    break                                # barge-in
                spoken.append(sentence)
        finally:
            producer.cancel()
            if speaking and not self.speaker.cancelled:
                await self.client.agent_idle()
        reply = " ".join(spoken) + (" [cut off]" if self.speaker.cancelled else "")
        self.history += [user, {"role": "assistant", "content": reply}]


# ── Audio sources ────────────────────────────────────────────────────────────

async def microphone_frames() -> AsyncIterator[bytes]:
    import sounddevice as sd

    loop = asyncio.get_running_loop()
    queue: asyncio.Queue = asyncio.Queue()

    def callback(indata, frames, time_info, status):  # audio thread
        loop.call_soon_threadsafe(queue.put_nowait, bytes(indata))

    with sd.RawInputStream(samplerate=SAMPLE_RATE, channels=1, dtype="int16",
                           blocksize=FRAME_SAMPLES, callback=callback):
        while True:
            yield await queue.get()


async def file_frames(paths: list[Path], gap_s: float) -> AsyncIterator[bytes]:
    """Stream WAVs in real time with `gap_s` of silence after each, then silence forever."""
    silence = bytes(FRAME_SAMPLES * 2)
    frame_s = FRAME_SAMPLES / SAMPLE_RATE
    clock = time.perf_counter()
    for path in paths:
        audio = read_wav_16k(path).tobytes()
        frames = [audio[i:i + FRAME_SAMPLES * 2] for i in range(0, len(audio), FRAME_SAMPLES * 2)]
        frames += [silence] * int(gap_s / frame_s)
        for frame in frames:
            clock += frame_s
            await asyncio.sleep(max(0.0, clock - time.perf_counter()))
            yield frame.ljust(FRAME_SAMPLES * 2, b"\0")
    while True:
        clock += frame_s
        await asyncio.sleep(max(0.0, clock - time.perf_counter()))
        yield silence


# ── Main ─────────────────────────────────────────────────────────────────────

async def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--url", default=os.environ.get("VSIP_URL", "wss://api.vsip.online/v1/stream"))
    parser.add_argument("--language", default=os.environ.get("LANGUAGE"), help="e.g. en (optional)")
    parser.add_argument("--enroll-seconds", type=float, default=18.0, help="live enrollment length")
    parser.add_argument("--enroll-file", type=Path, help="enroll from a WAV instead of the mic")
    parser.add_argument("--input-file", type=Path, action="append", help="stream WAV(s) instead of the mic")
    parser.add_argument("--gap", type=float, default=8.0, help="silence after each --input-file")
    parser.add_argument("--no-playback", action="store_true", help="don't open the speakers")
    parser.add_argument("--save-replies", type=Path, help="write each spoken reply as a WAV here")
    parser.add_argument("--duration", type=float, default=0, help="stop after N seconds (0 = run forever)")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S")
    for noisy in ("httpx", "vsip.client", "websockets"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if not os.environ.get("VSIP_API_KEY") or not os.environ.get("OPENAI_API_KEY"):
        print("Set VSIP_API_KEY and OPENAI_API_KEY (see .env.example).")
        return 1

    client = VSIPClient(api_key=os.environ["VSIP_API_KEY"], url=args.url, mode="full")
    speaker = Speaker(enabled=not args.no_playback, save_dir=args.save_replies)
    agent = Agent(client, AsyncOpenAI(), speaker, args.language)
    session = VoiceSession(client, on_turn=agent.on_turn, on_barge_in=agent.on_barge_in, on_event=agent.on_event)
    agent.session = session

    use_files = bool(args.input_file or args.enroll_file)
    async with client:
        await session.start()
        log.info("connected · stream %s", client.stream_id)

        # 1. Enrollment: 15-30 s of natural speech, locked atomically.
        if args.enroll_file:
            enroll_pcm = read_wav_16k(args.enroll_file).tobytes()
        else:
            print(f"\nEnrollment — read this aloud for about {args.enroll_seconds:.0f} seconds:\n\n  {ENROLL_PROMPT}\n")
            captured, needed = bytearray(), int(args.enroll_seconds * SAMPLE_RATE * 2)
            async for frame in microphone_frames():
                captured += frame
                await session.send_audio(frame)
                if len(captured) >= needed:
                    break
            enroll_pcm = bytes(captured)
        await client.enroll_speaker("primary", int16_to_float32_bytes(enroll_pcm), lock=True)

        # 2. Converse: every frame goes to VSIP; VoiceSession calls on_turn.
        source = file_frames(args.input_file or [], args.gap) if use_files else microphone_frames()
        if not use_files:
            print("Ready — talk to the agent. Interrupt it any time. Ctrl+C to quit.\n")
        deadline = time.monotonic() + args.duration if args.duration else None
        async for frame in source:
            await session.send_audio(frame)
            if deadline and time.monotonic() > deadline:
                break
        await session.stop()
    return 0


if __name__ == "__main__":
    try:
        sys.exit(asyncio.run(main()))
    except KeyboardInterrupt:
        pass