#!/usr/bin/env python3
"""Runtime patches for the vllm-omni Qwen3-TTS container.

- Propagate WebSocket session.config.seed into every generated speech segment.
- Log WebSocket config and segment generation on the real 8091 endpoint.
"""
from pathlib import Path

STREAM = Path('/usr/local/lib/python3.12/dist-packages/vllm_omni/entrypoints/openai/serving_speech_stream.py')
AUDIO = Path('/usr/local/lib/python3.12/dist-packages/vllm_omni/entrypoints/openai/protocol/audio.py')


def replace_once(text: str, old: str, new: str, label: str) -> str:
    if new in text:
        return text
    if old not in text:
        raise RuntimeError(f'Patch anchor not found: {label}')
    return text.replace(old, new, 1)


def patch_seed_field() -> None:
    """Add an optional `seed` field to StreamingSpeechSessionConfig.

    The streaming WebSocket config model in protocol/audio.py ships without a
    `seed` attribute, but the stream server reads `config.seed` (see seed
    propagation patch below). Without this field every session.config fails with
    AttributeError: 'StreamingSpeechSessionConfig' object has no attribute 'seed'.
    """
    text = AUDIO.read_text()
    # Anchor on this class's unique docstring + first field: the task_type line
    # is not unique (it also appears in SpeechBatchItem / BatchSpeechRequest).
    new = replace_once(
        text,
        '    """Configuration sent as the first WebSocket message for streaming TTS."""\n\n'
        '    model: str | None = None\n',
        '    """Configuration sent as the first WebSocket message for streaming TTS."""\n\n'
        '    seed: int | None = None\n'
        '    model: str | None = None\n',
        'StreamingSpeechSessionConfig.seed field',
    )
    if new != text:
        AUDIO.write_text(new)
        print('[runtime-patch] StreamingSpeechSessionConfig.seed field added', flush=True)
    else:
        print('[runtime-patch] StreamingSpeechSessionConfig.seed field already present', flush=True)


def main() -> None:
    patch_seed_field()

    text = STREAM.read_text()
    original = text

    text = replace_once(
        text,
        '            if config.model and hasattr(self._speech_service, "_check_model"):\n',
        '            logger.info(\n'
        '                "Streaming speech config: model=%s voice=%s task_type=%s language=%s seed=%s split=%s stream_audio=%s instructions=%r speaker_embedding_dim=%s",\n'
        '                config.model,\n'
        '                config.voice,\n'
        '                config.task_type,\n'
        '                config.language,\n'
        '                config.seed,\n'
        '                config.split_granularity,\n'
        '                config.stream_audio,\n'
        '                config.instructions,\n'
        '                len(config.speaker_embedding) if config.speaker_embedding is not None else None,\n'
        '            )\n\n'
        '            if config.model and hasattr(self._speech_service, "_check_model"):\n',
        'config logging',
    )

    text = replace_once(
        text,
        '            stream=config.stream_audio,\n        )\n',
        '            stream=config.stream_audio,\n            seed=config.seed,\n        )\n'
        '        logger.info(\n'
        '            "Streaming speech segment: index=%s chars=%s voice=%s task_type=%s language=%s seed=%s speaker_embedding_dim=%s",\n'
        '            sentence_index,\n'
        '            len(sentence_text),\n'
        '            config.voice,\n'
        '            config.task_type,\n'
        '            config.language,\n'
        '            config.seed,\n'
        '            len(config.speaker_embedding) if config.speaker_embedding is not None else None,\n'
        '        )\n',
        'seed propagation and segment logging',
    )

    # --- Silence padding (PCM streaming only) ---
    # Lead-in silence at session start gives the client time to set up smooth
    # playback before real audio (mitigates start-of-stream stutter). Trailing
    # silence avoids the abrupt cut-off at the end. Both in ms via env vars.
    text = replace_once(
        text,
        '_PCM_SAMPLE_RATE = 24000\n',
        '_PCM_SAMPLE_RATE = 24000\n'
        'import os as _os\n'
        '_LEADIN_SILENCE_MS = int(_os.environ.get("QWEN3_TTS_LEADIN_SILENCE_MS", "250"))\n'
        '_TRAILING_SILENCE_MS = int(_os.environ.get("QWEN3_TTS_TRAILING_SILENCE_MS", "300"))\n'
        '_LEADIN_SILENCE_BYTES = b"\\x00" * (int(_PCM_SAMPLE_RATE * _LEADIN_SILENCE_MS / 1000) * 2)\n'
        '_TRAILING_SILENCE_BYTES = b"\\x00" * (int(_PCM_SAMPLE_RATE * _TRAILING_SILENCE_MS / 1000) * 2)\n',
        'silence padding constants',
    )

    text = replace_once(
        text,
        '        await websocket.send_json(start_payload)\n',
        '        await websocket.send_json(start_payload)\n'
        '        if config.stream_audio and response_format == "pcm" and sentence_index == 0 and _LEADIN_SILENCE_BYTES:\n'
        '            await websocket.send_bytes(_LEADIN_SILENCE_BYTES)\n',
        'lead-in silence',
    )

    text = replace_once(
        text,
        '                    # Send session.done\n'
        '                    await websocket.send_json(\n',
        '                    if config.stream_audio and config.response_format == "pcm" and _TRAILING_SILENCE_BYTES:\n'
        '                        await websocket.send_bytes(_TRAILING_SILENCE_BYTES)\n\n'
        '                    # Send session.done\n'
        '                    await websocket.send_json(\n',
        'trailing silence',
    )

    if text != original:
        STREAM.write_text(text)
        print('[runtime-patch] WebSocket seed propagation/logging + silence padding applied', flush=True)
    else:
        print('[runtime-patch] WebSocket seed propagation/logging already present', flush=True)


if __name__ == '__main__':
    main()