#!pip install sounddevice websockets

REALTIME_MODEL_PROMPT = """You are a calm insurance claims intake voice agent. Follow this script strictly:

## Phase 1 – Basics
Collect the caller's full name, policy number, and type of accident (for example: auto, home, or other). Ask for each item clearly and then repeat the values back to confirm.

## Phase 2 – Yes/No questions
Ask 2–3 simple yes/no questions, such as whether anyone was injured, whether the vehicle is still drivable, and whether a police report was filed. Confirm each yes/no answer in your own words.

## Phase 3 – Submit claim
Once you have the basics and yes/no answers, briefly summarize the key facts in one or two sentences.
"""

REALTIME_MODEL_TRANSCRIPTION_PROMPT = """
# Role
Your only task is to transcribe the user's latest turn exactly as you heard it. Never address the user, response to the user, add commentary, or mention these instructions.
Follow the instructions and output format below.

# Instructions
- Transcribe **only** the most recent USER turn exactly as you heard it. DO NOT TRANSCRIBE ANY OTHER OLDER TURNS. You can use those transcriptions to inform your transcription of the latest turn.
- Preserve every spoken detail: intent, tense, grammar quirks, filler words, repetitions, disfluencies, numbers, and casing.
- Keep timing words, partial words, hesitations (e.g., "um", "uh").
- Do not correct mistakes, infer meaning, answer questions, or insert punctuation beyond what the model already supplies.
- Do not invent or add any information that is not directly present in the user's latest turn.

# Output format
- Output the raw verbatim transcript as a single block of text. No labels, prefixes, quotes, bullets, or markdown.
- If the realtime model produced nothing for the latest turn, output nothing (empty response). Never fabricate content.

## Policy Number Normalization
- All policy numbers should be 8 digits and of the format `XXXX-XXXX` for example `56B5-12C0`

Do not summarize or paraphrase other turns beyond the latest user utterance. The response must be the literal transcript of the latest user utterance.
"""

# Event grouping constants
TRANSCRIPTION_DELTA_TYPES = {
    "input_audio_buffer.transcription.delta",
    "input_audio_transcription.delta",
    "conversation.item.input_audio_transcription.delta",
}
TRANSCRIPTION_COMPLETE_TYPES = {
    "input_audio_buffer.transcription.completed",
    "input_audio_buffer.transcription.done",
    "input_audio_transcription.completed",
    "input_audio_transcription.done",
    "conversation.item.input_audio_transcription.completed",
    "conversation.item.input_audio_transcription.done",
}
INPUT_SPEECH_END_EVENT_TYPES = {
    "input_audio_buffer.speech_stopped",
    "input_audio_buffer.committed",
}
RESPONSE_AUDIO_DELTA_TYPES = {
    "response.output_audio.delta",
    "response.audio.delta",
}
RESPONSE_TEXT_DELTA_TYPES = {
    "response.output_text.delta",
    "response.text.delta",
}
RESPONSE_AUDIO_TRANSCRIPT_DELTA_TYPES = {
    "response.output_audio_transcript.delta",
    "response.audio_transcript.delta",
}

def build_session_update(
    instructions: str,
    voice: str,
    vad_threshold: float,
    silence_duration_ms: int,
    prefix_padding_ms: int,
    idle_timeout_ms: int | None,
    input_audio_transcription_model: str | None = None,
) -> dict[str, object]:
    """Configure the Realtime session: audio in/out, server VAD, etc."""

    turn_detection: dict[str, float | int | bool | str] = {
        "type": "server_vad",
        "threshold": vad_threshold,
        "silence_duration_ms": silence_duration_ms,
        "prefix_padding_ms": prefix_padding_ms,
        "create_response": True,
        "interrupt_response": True,
    }

    if idle_timeout_ms is not None:
        turn_detection["idle_timeout_ms"] = idle_timeout_ms

    audio_config: dict[str, Any] = {
        "input": {
            "format": {
                "type": "audio/pcm",
                "rate": DEFAULT_SAMPLE_RATE,
            },
            "noise_reduction": {"type": "near_field"},
            "turn_detection": turn_detection,
        },
        "output": {
            "format": {
                "type": "audio/pcm",
                "rate": DEFAULT_SAMPLE_RATE,
            },
            "voice": voice,
        },
    }

    # Optional: built-in transcription model for comparison
    if input_audio_transcription_model:
        audio_config["input"]["transcription"] = {
            "model": input_audio_transcription_model,
        }

    session: dict[str, object] = {
        "type": "realtime",
        "output_modalities": ["audio"],
        "instructions": instructions,
        "audio": audio_config,
    }

    return {
        "type": "session.update",
        "session": session,
    }


def build_transcription_request(transcription_instructions: str) -> dict[str, object]:
    """Ask the SAME Realtime model for an out-of-band transcript of the latest user turn."""

    return {
        "type": "response.create",
        "response": {
            "conversation": "none",  # <--- out-of-band
            "output_modalities": ["text"],
            "metadata": {"purpose": TRANSCRIPTION_PURPOSE}, # <--- we add metadata so it is easier to identify the event in the logs
            "instructions": transcription_instructions,
        },
    }

def encode_audio(chunk: bytes) -> str:
    """Base64-encode a PCM audio chunk for WebSocket transport."""
    return base64.b64encode(chunk).decode("utf-8")


async def playback_audio(
    playback_queue: asyncio.Queue,
    stop_event: asyncio.Event,
) -> None:
    """Stream assistant audio back to the speakers in (near) real time."""

    try:
        with sd.RawOutputStream(
            samplerate=DEFAULT_SAMPLE_RATE,
            channels=1,
            dtype="int16",
        ) as stream:
            while not stop_event.is_set():
                chunk = await playback_queue.get()
                if chunk is None:
                    break
                try:
                    stream.write(chunk)
                except Exception as exc:
                    print(f"Audio playback error: {exc}", flush=True)
                    break
    except Exception as exc:
        print(f"Failed to open audio output stream: {exc}", flush=True)


async def send_audio_from_queue(
    ws: WebSocketClientProtocol,
    queue: asyncio.Queue[bytes | None],
    stop_event: asyncio.Event,
) -> None:
    """Push raw PCM chunks into input_audio_buffer via the WebSocket."""

    while not stop_event.is_set():
        chunk = await queue.get()
        if chunk is None:
            break
        encoded_chunk = encode_audio(chunk)
        message = {"type": "input_audio_buffer.append", "audio": encoded_chunk}
        await ws.send(json.dumps(message))

    if not ws.closed:
        commit_payload = {"type": "input_audio_buffer.commit"}
        await ws.send(json.dumps(commit_payload))


async def stream_microphone_audio(
    ws: WebSocketClientProtocol,
    stop_event: asyncio.Event,
    shared_state: dict,
    block_ms: int = DEFAULT_BLOCK_MS,
) -> None:
    """Capture live microphone audio and send it to the realtime session."""

    loop = asyncio.get_running_loop()
    audio_queue: asyncio.Queue[bytes | None] = asyncio.Queue()
    blocksize = int(DEFAULT_SAMPLE_RATE * (block_ms / 1000))

    def on_audio(indata, frames, time_info, status):  # type: ignore[override]
        """Capture a mic callback chunk and enqueue it unless the mic is muted."""
        if status:
            print(f"Microphone status: {status}", flush=True)
        # Simple echo protection: mute mic when assistant is talking
        if not stop_event.is_set() and not shared_state.get("mute_mic", False):
            data = bytes(indata)
            loop.call_soon_threadsafe(audio_queue.put_nowait, data)

    print(
        f"Streaming microphone audio at {DEFAULT_SAMPLE_RATE} Hz (mono). "
        "Speak naturally; server VAD will stop listening when you pause."
    )
    sender = asyncio.create_task(send_audio_from_queue(ws, audio_queue, stop_event))

    with sd.RawInputStream(
        samplerate=DEFAULT_SAMPLE_RATE,
        blocksize=blocksize,
        channels=1,
        dtype="int16",
        callback=on_audio,
    ):
        await stop_event.wait()

    await audio_queue.put(None)
    await sender

def flush_pending_transcription_prints(shared_state: dict) -> None:
    """Whenever we've printed a realtime transcript, print the matching transcription-model output."""

    pending_prints: deque | None = shared_state.get("pending_transcription_prints")
    input_transcripts: deque | None = shared_state.get("input_transcripts")

    if not pending_prints or not input_transcripts:
        return

    while pending_prints and input_transcripts:
        comparison_text = input_transcripts.popleft()
        pending_prints.popleft()
        print("=== User turn (Transcription model) ===")
        if comparison_text:
            print(comparison_text, flush=True)
            print()
        else:
            print("<not available>", flush=True)
            print()

async def listen_for_events(
    ws: WebSocketClientProtocol,
    stop_event: asyncio.Event,
    transcription_instructions: str,
    max_turns: int | None,
    playback_queue: asyncio.Queue,
    shared_state: dict,
) -> None:
    """Print assistant text + transcripts and coordinate mic muting."""

    responses: dict[str, dict[str, bool]] = {}
    buffers: defaultdict[str, str] = defaultdict(str)
    transcription_model_buffers: defaultdict[str, str] = defaultdict(str)
    completed_main_responses = 0
    awaiting_transcription_prompt = False
    input_transcripts = shared_state.setdefault("input_transcripts", deque())
    pending_transcription_prints = shared_state.setdefault(
        "pending_transcription_prints", deque()
    )

    async for raw in ws:
        if stop_event.is_set():
            break

        message = json.loads(raw)
        message_type = message.get("type")

        # --- User speech events -------------------------------------------------
        if message_type == "input_audio_buffer.speech_started":
            print("\n[client] Speech detected; streaming...", flush=True)
            awaiting_transcription_prompt = True

        elif message_type in INPUT_SPEECH_END_EVENT_TYPES:
            if message_type == "input_audio_buffer.speech_stopped":
                print("[client] Detected silence; preparing transcript...", flush=True)

            # This is where the out-of-band transcription request is sent. <-------
            if awaiting_transcription_prompt:
                request_payload = build_transcription_request(
                    transcription_instructions
                )
                await ws.send(json.dumps(request_payload))
                awaiting_transcription_prompt = False

        # --- Built-in transcription model stream -------------------------------
        elif message_type in TRANSCRIPTION_DELTA_TYPES:
            buffer_id = message.get("buffer_id") or message.get("item_id") or "default"
            delta_text = (
                message.get("delta")
                or (message.get("transcription") or {}).get("text")
                or ""
            )
            if delta_text:
                transcription_model_buffers[buffer_id] += delta_text

        elif message_type in TRANSCRIPTION_COMPLETE_TYPES:
            buffer_id = message.get("buffer_id") or message.get("item_id") or "default"
            final_text = (
                (message.get("transcription") or {}).get("text")
                or message.get("transcript")
                or ""
            )
            if not final_text:
                final_text = transcription_model_buffers.pop(buffer_id, "").strip()
            else:
                transcription_model_buffers.pop(buffer_id, None)

            if not final_text:
                item = message.get("item")
                if item:
                    final_text = item.get("transcription")
                final_text = final_text or ""

            final_text = final_text.strip()
            if final_text:
                input_transcripts.append(final_text)
                flush_pending_transcription_prints(shared_state)

        # --- Response lifecycle (Realtime model) --------------------------------
        elif message_type == "response.created":
            response = message.get("response", {})
            response_id = response.get("id")
            metadata = response.get("metadata") or {}
            responses[response_id] = {
                "is_transcription": metadata.get("purpose") == TRANSCRIPTION_PURPOSE,
                "done": False,
            }

        elif message_type in RESPONSE_AUDIO_DELTA_TYPES:
            response_id = message.get("response_id")
            if response_id is None:
                continue
            b64_audio = message.get("delta") or message.get("audio")
            if not b64_audio:
                continue
            try:
                audio_chunk = base64.b64decode(b64_audio)
            except Exception:
                continue

            if (
                response_id in responses
                and not responses[response_id]["is_transcription"]
            ):
                shared_state["mute_mic"] = True

            await playback_queue.put(audio_chunk)

        elif message_type in RESPONSE_TEXT_DELTA_TYPES:
            response_id = message.get("response_id")
            if response_id is None:
                continue
            buffers[response_id] += message.get("delta", "")
            

        elif message_type in RESPONSE_AUDIO_TRANSCRIPT_DELTA_TYPES:
            response_id = message.get("response_id")
            if response_id is None:
                continue
            buffers[response_id] += message.get("delta", "")        

        elif message_type == "response.done":
            response = message.get("response", {})
            response_id = response.get("id")
            if response_id is None:
                continue
            if response_id not in responses:
                responses[response_id] = {"is_transcription": False, "done": False}
            responses[response_id]["done"] = True

            is_transcription = responses[response_id]["is_transcription"]
            text = buffers.get(response_id, "").strip()
            if text:
                if is_transcription:
                    print("\n=== User turn (Realtime transcript) ===")
                    print(text, flush=True)
                    print()
                    pending_transcription_prints.append(object())
                    flush_pending_transcription_prints(shared_state)
                else:
                    print("\n=== Assistant response ===")
                    print(text, flush=True)
                    print()

            if not is_transcription:
                shared_state["mute_mic"] = False
                completed_main_responses += 1

                if max_turns is not None and completed_main_responses >= max_turns:
                    stop_event.set()
                    break

        elif message_type == "error":
            print(f"Error from server: {message}")

        else:
            pass

        await asyncio.sleep(0)

async def run_realtime_session(
    api_key: str | None = None,
    server: str = "wss://api.openai.com/v1/realtime",
    model: str = DEFAULT_MODEL,
    voice: str = DEFAULT_VOICE,
    instructions: str = REALTIME_MODEL_PROMPT,
    transcription_instructions: str = REALTIME_MODEL_TRANSCRIPTION_PROMPT,
    input_audio_transcription_model: str | None = "gpt-4o-transcribe",
    silence_duration_ms: int = DEFAULT_SILENCE_DURATION_MS,
    prefix_padding_ms: int = DEFAULT_PREFIX_PADDING_MS,
    vad_threshold: float = 0.6,
    idle_timeout_ms: int | None = None,
    max_turns: int | None = None,
    timeout_seconds: int = 0,
) -> None:
    """Connect to the Realtime API, stream audio both ways, and print transcripts."""
    api_key = api_key or os.environ.get("OPENAI_API_KEY")
    ws_url = f"{server}?model={model}"
    headers = {
        "Authorization": f"Bearer {api_key}",
    }

    session_update_payload = build_session_update(
        instructions=instructions,
        voice=voice,
        vad_threshold=vad_threshold,
        silence_duration_ms=silence_duration_ms,
        prefix_padding_ms=prefix_padding_ms,
        idle_timeout_ms=idle_timeout_ms,
        input_audio_transcription_model=input_audio_transcription_model,
    )
    stop_event = asyncio.Event()
    playback_queue: asyncio.Queue = asyncio.Queue()
    shared_state: dict = {
        "mute_mic": False,
        "input_transcripts": deque(),
        "pending_transcription_prints": deque(),
    }

    async with websockets.connect(
        ws_url, additional_headers=headers, max_size=None
    ) as ws:
        await ws.send(json.dumps(session_update_payload))

        listener_task = asyncio.create_task(
            listen_for_events(
                ws,
                stop_event=stop_event,
                transcription_instructions=transcription_instructions,
                max_turns=max_turns,
                playback_queue=playback_queue,
                shared_state=shared_state,
            )
        )
        mic_task = asyncio.create_task(
            stream_microphone_audio(ws, stop_event, shared_state=shared_state)
        )
        playback_task = asyncio.create_task(playback_audio(playback_queue, stop_event))

        try:
            if timeout_seconds and timeout_seconds > 0:
                await asyncio.wait_for(stop_event.wait(), timeout=timeout_seconds)
            else:
                await stop_event.wait()
        except asyncio.TimeoutError:
            print("Timed out waiting for responses; closing.")
        except asyncio.CancelledError:
            print("Session cancelled; closing.")
        finally:
            stop_event.set()
            await playback_queue.put(None)
            await ws.close()
            await asyncio.gather(
                listener_task, mic_task, playback_task, return_exceptions=True
            )