#!/usr/bin/env python3
"""Самостоятельный разговор WebRTC V3 по контракту Codex CLI 0.156.1.

Не запускает Codex/Electron/app-server/helper. Не выполняет команды модели.
Приватный ChatGPT API экспериментален: отказ не заменяется другим транспортом.
"""
import argparse
import asyncio
from collections import deque
from fractions import Fraction
import json
import math
import os
from pathlib import Path
import re
import sys
import tempfile
import uuid
import wave

import aiohttp
import av
from aiortc import AudioStreamTrack, RTCConfiguration, RTCPeerConnection
from aiortc import RTCRtpSender, RTCSessionDescription
from aiortc.mediastreams import MediaStreamError

# Это именно backend V3 signaling и отдельный sideband, не GA /v1/realtime.
CALL_URL = "https://chatgpt.com/backend-api/codex/realtime/calls?intent=quicksilver&architecture=avas"
SIDEBAND = "wss://api.openai.com/v1/live/"
VOICES = ("juniper", "maple", "spruce", "ember", "vale", "breeze", "arbor", "sol", "cove")
DEFAULT_MODEL = "gpt-live-1-codex"  # core/realtime_conversation.rs, CLI 0.156.1 V3.
RATE, SAMPLES = 48000, 960  # 20 мс PCM16 mono; RTP Opus имеет clock 48 кГц.
FRAME_BYTES = SAMPLES * 2
SILENCE = bytes(FRAME_BYTES)
INSTRUCTIONS = ("Разговаривай по-русски. Отвечай коротко на услышанную речь. "
                "Не делегируй работу, не вызывай инструменты, не выполняй команды. "
                "Если просят действие над компьютером, объясни, что доступен только разговор.")


class VoiceError(RuntimeError):
    """Только контролируемый текст: никаких тел сервера, токенов и SDP в ошибке."""


def read_auth(path):
    # Основной профиль только читается. Refresh и запись auth.json запрещены.
    auth = json.loads(path.expanduser().read_text(encoding="utf-8"))
    tokens = auth.get("tokens") or {}
    token = tokens.get("access_token")
    if auth.get("auth_mode") != "chatgpt" or not isinstance(token, str) or not token or token == "0000" or any(c in token for c in "\r\n"):
        raise VoiceError("auth: нужен собственный ChatGPT access_token; 0000 не подходит")
    headers = {"Authorization": "Bearer " + token, "User-Agent": "codex-voice-isolated/0.1 aiortc/1.15.0",
               # Обычный version negotiation из realtime_request_headers, не
               # attestation и не подмена браузера. У V3 здесь quicksilver=v2.
               "OpenAI-Alpha": "quicksilver=v2", "x-session-id": str(uuid.uuid4()),
               "originator": "codex_voice_isolated"}
    account = tokens.get("account_id")
    if account:
        if not isinstance(account, str) or any(c in account for c in "\r\n"):
            raise VoiceError("auth: некорректный account_id")
        headers["ChatGPT-Account-ID"] = account
    return headers


def call_body(sdp, voice, model=None):
    # Session задаётся уже при создании call: inference может начаться сразу.
    if voice not in VOICES:
        raise VoiceError("voice: неизвестный голос каталога V3 / CLI 0.156.1")
    session = {"instructions": INSTRUCTIONS, "audio": {"output": {"voice": voice}},
               "delegation": {"type": "client"}}
    # В RPC model необязателен, но ядро CLI подставляет default перед HTTP.
    # Поэтому отсутствие поля в TUI НЕ означает его отсутствие в call JSON.
    session["model"] = model or DEFAULT_MODEL
    return {"sdp": sdp, "session": session}


def no_redirect_trace():
    # aiohttp ws_connect не принимает allow_redirects=False. Trace прерывает
    # WSS redirect до следующего запроса, чтобы Bearer не сменил назначение.
    trace = aiohttp.TraceConfig()
    async def reject(session, context, params):
        raise VoiceError("network: перенаправление запрещено")
    trace.on_request_redirect.append(reject)
    return trace


def call_id_from_location(location):
    # Location не используем как адрес для отправки Bearer на произвольный хост.
    # Берём только последний безопасный сегмент, WSS-хост зафиксирован выше.
    from urllib.parse import urlsplit
    path = urlsplit(location or "").path.rstrip("/")
    call_id = path.rsplit("/", 1)[-1]
    if not re.fullmatch(r"[A-Za-z0-9_-]{1,200}", call_id):
        raise VoiceError("signaling: нет безопасного call ID в Location")
    return call_id


def read_wav(path):
    # Сами WAV/RIFF-байты в WebRTC не идут: track получает только PCM-сэмплы.
    with wave.open(str(path.expanduser()), "rb") as audio:
        if (audio.getframerate(), audio.getnchannels(), audio.getsampwidth(), audio.getcomptype()) != (RATE, 1, 2, "NONE"):
            raise VoiceError("file: нужен WAV PCM16 mono 48000 Гц; конвертируйте ffmpeg")
        frames = audio.getnframes()
        if not RATE // 5 <= frames <= RATE * 60:
            raise VoiceError("file: длительность должна быть 0,2–60 секунд")
        pcm = audio.readframes(frames)
        if len(pcm) != frames * 2:
            raise VoiceError("file: WAV обрезан")
        return pcm


class InputTrack(AudioStreamTrack):
    """Тактируемый track: микрофон/файл или синтетическая тишина, без backlog."""
    def __init__(self, pcm=None, device=None):
        super().__init__()
        self.pcm, self.device = pcm, device
        self.offset = 0
        self.queue = deque(maxlen=10)  # Не накапливаем секунды запоздавшего голоса.
        self.pending = bytearray()
        self.capture = None
        self.muted = True
        self.epoch = 0
        self.timestamp = 0
        self.clock = None
        self.file_done = asyncio.Event()
        self.loop = asyncio.get_running_loop()

    def enqueue(self, data, epoch):
        # Callback старого устройства может уже стоять в очереди asyncio.
        if self.muted or epoch != self.epoch:
            return
        self.pending.extend(data)
        while len(self.pending) >= FRAME_BYTES:
            self.queue.append(bytes(self.pending[:FRAME_BYTES]))
            del self.pending[:FRAME_BYTES]

    def unmute(self):
        if not self.muted:
            return
        self.epoch += 1
        epoch = self.epoch
        self.muted = False
        if self.pcm is None:
            import sounddevice as sd
            # Событие записи начинается только после /voice и соединения.
            def callback(data, frames, time_info, status):
                self.loop.call_soon_threadsafe(self.enqueue, bytes(data), epoch)
            try:
                self.capture = sd.RawInputStream(samplerate=RATE, channels=1, dtype="int16",
                                                 blocksize=SAMPLES, device=self.device, callback=callback)
                self.capture.start()
            except Exception:
                self.mute()
                raise VoiceError("microphone: не удалось открыть выбранное устройство") from None

    def mute(self):
        # Прекращаем реальный захват; track продолжает отправлять нулевые сэмплы.
        self.muted = True
        self.epoch += 1
        self.queue.clear()
        self.pending.clear()
        if self.capture is not None:
            capture, self.capture = self.capture, None
            failed = False
            try:
                capture.abort()
            except Exception:
                failed = True
            try:
                capture.close()
            except Exception:
                failed = True
            if failed:
                raise VoiceError("microphone: ошибка закрытия устройства; захват в track отключён")

    def next_pcm(self):
        if self.muted:
            return SILENCE
        if self.pcm is not None:
            data = self.pcm[self.offset:self.offset + FRAME_BYTES]
            self.offset += len(data)
            if self.offset >= len(self.pcm):
                self.file_done.set()
            return data.ljust(FRAME_BYTES, b"\0")
        return self.queue.popleft() if self.queue else SILENCE

    async def recv(self):
        if self.readyState != "live":
            raise MediaStreamError
        if self.clock is None:
            self.clock = self.loop.time()
        await asyncio.sleep(max(0, self.clock + self.timestamp / RATE - self.loop.time()))
        frame = av.AudioFrame(format="s16", layout="mono", samples=SAMPLES)
        frame.planes[0].update(self.next_pcm())
        frame.sample_rate, frame.time_base, frame.pts = RATE, Fraction(1, RATE), self.timestamp
        self.timestamp += SAMPLES
        return frame

    def stop(self):
        try:
            self.mute()
        finally:
            super().stop()


def caption(event):
    # Имена/поля V3 взяты из protocol_frameless_bidi.rs, не из GA V2.
    kind = event.get("type")
    if kind in ("input_transcript.added", "output_transcript.added"):
        text = (event.get("item") or {}).get("text")
        role = "user" if kind.startswith("input") else "assistant"
        return (role, False, text) if isinstance(text, str) else None
    if kind == "turn.done":
        turn = event.get("turn") or {}
        role, text = turn.get("role"), turn.get("transcript")
        return (role, True, text) if role in ("user", "assistant") and isinstance(text, str) else None
    return None


class VoiceSession:
    def __init__(self, args, pcm=None, emit=print):
        self.args, self.emit = args, emit
        self.track = InputTrack(pcm, args.input_device)
        self.pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
        self.http = self.ws = self.speaker = None
        self.tasks = []
        self.connected, self.finished = asyncio.Event(), asyncio.Event()
        self.failed = None
        self.closed = False
        self.output = bytearray()
        self.received_samples = self.nonzero_samples = 0
        self.transcript_chars = {"user": 0, "assistant": 0}
        self.assistant_done = asyncio.Event()
        self.phase = "initial"

    def fail(self, message):
        self.failed = VoiceError(message)
        # Даже если connection/sideband сломался до пробуждения консоли,
        # реальный захват и очередь больше не должны продолжаться.
        try:
            self.track.mute()
        except VoiceError:
            self.emit("microphone: устройство сообщило ошибку закрытия")
        self.finished.set()

    async def start(self):
        try:
            await self._start()
        except BaseException:
            await self.close()
            raise

    async def _start(self):
        headers = read_auth(self.args.auth_file)
        self.http = self.args.network.client(headers, no_redirect_trace())
        transceiver = self.pc.addTransceiver(self.track, direction="sendrecv")
        # Только Opus; SDP channels=2 не означает, что входной файл stereo.
        transceiver.setCodecPreferences([c for c in RTCRtpSender.getCapabilities("audio").codecs
                                         if c.mimeType.lower() == "audio/opus"])
        channel = self.pc.createDataChannel("oai-events")
        channel.on("message", lambda message: None)  # Как helper: captions идут по sideband.

        @self.pc.on("connectionstatechange")
        async def state_change():
            if self.pc.connectionState == "connected":
                self.connected.set()
            elif self.pc.connectionState == "failed":
                self.fail("media: WebRTC connection failed")

        @self.pc.on("track")
        def remote_track(track):
            if track.kind == "audio":
                self.tasks.append(asyncio.create_task(self.receive_audio(track)))

        self.phase = "offer"
        await self.pc.setLocalDescription(await self.pc.createOffer())
        self.phase = "signaling"
        body = call_body(self.pc.localDescription.sdp, self.args.voice, self.args.model)
        async with self.http.post(CALL_URL, json=body, allow_redirects=False) as response:
            if response.status != 200 and response.status != 201:
                challenge = response.headers.get("cf-mitigated") == "challenge"
                # Для диагностики выводим только машинный code/param. Серверные
                # message/SDP/идентификаторы/заголовки не попадают в журнал.
                code = param = "unknown"
                try:
                    rejected = json.loads(await response.content.read(8192))
                    error = rejected.get("error") or {}
                    for key in ("code", "param"):
                        value = error.get(key)
                        if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_.-]{1,80}", value):
                            if key == "code":
                                code = value
                            else:
                                param = value
                except (ValueError, AttributeError, TypeError):
                    pass
                raise VoiceError(f"signaling: HTTP {response.status}; challenge={challenge}; code={code}; param={param}; тело скрыто")
            # Не выводим ответ: SDP содержит IP/порты/ICE credentials.
            raw = await response.content.read(200001)
            if len(raw) > 200000:
                raise VoiceError("signaling: слишком большой SDP answer")
            answer = raw.decode("utf-8")  # CLI тоже принимает plain SDP, не JSON.
            if not answer.startswith("v=0"):
                raise VoiceError("signaling: ответ не является SDP")
            call_id = call_id_from_location(response.headers.get("Location"))
        self.phase = "answer"
        await self.pc.setRemoteDescription(RTCSessionDescription(sdp=answer, type="answer"))
        self.phase = "sideband"
        # V3 joins existing call: никакого повторного session.update!
        self.ws = await self.http.ws_connect(SIDEBAND + call_id, heartbeat=20, max_msg_size=1000000)
        self.tasks.append(asyncio.create_task(self.receive_events()))
        self.phase = "media"
        await asyncio.wait_for(self.connected.wait(), timeout=self.args.timeout)
        if self.failed:
            raise self.failed
        self.track.unmute()
        self.emit("WebRTC V3 соединён; микрофон включён" if self.track.pcm is None else "WebRTC V3 соединён; передаётся файл")

    async def receive_events(self):
        try:
            async for message in self.ws:
                if message.type != aiohttp.WSMsgType.TEXT:
                    if message.type == aiohttp.WSMsgType.ERROR:
                        raise VoiceError("sideband: ошибка WSS")
                    continue
                event = json.loads(message.data)
                if not isinstance(event, dict):
                    raise VoiceError("sideband: неверный JSON event")
                kind = event.get("type")
                if kind == "delegation.created":
                    # Никогда не выполняем ни текст делегации, ни команды/tools.
                    raise VoiceError("delegation: сервер запросил агентную работу; не поддерживается, сеанс завершён")
                if kind == "error":
                    code = (event.get("error") or {}).get("code", "unknown")
                    safe_code = code if isinstance(code, str) and re.fullmatch(r"[a-zA-Z0-9_-]{1,80}", code) else "hidden"
                    raise VoiceError("sideband: service error " + safe_code)
                result = caption(event)
                if result:
                    role, final, text = result
                    if final:
                        self.transcript_chars[role] += len(text.strip())
                        # Стартовое приветствие не считается ответом на файл:
                        # сначала нужен подтверждённый user turn этого call.
                        if role == "assistant" and text.strip() and self.transcript_chars["user"]:
                            self.assistant_done.set()
                    # Файловая проба не публикует личную транскрипцию.
                    if self.track.pcm is None:
                        label = "Вы" if role == "user" else "Ассистент"
                        self.emit(f"{label} {'✓' if final else '…'} {text}")
            if not self.closed:
                self.fail("sideband: WSS завершён до остановки клиента")
        except asyncio.CancelledError:
            raise
        except Exception as error:
            self.fail(str(error) if isinstance(error, VoiceError) else "sideband: некорректное событие или сбой сети")

    async def receive_audio(self, track):
        # Декодированные RTP Opus сэмплы приводятся к формату устройства/WAV.
        resampler = av.AudioResampler(format="s16", layout="mono", rate=RATE)
        try:
            while True:
                for frame in resampler.resample(await track.recv()):
                    data = bytes(frame.planes[0])[:frame.samples * 2]
                    self.received_samples += frame.samples
                    self.nonzero_samples += sum(v != 0 for v in memoryview(data).cast("h"))
                    if self.args.output:
                        if len(self.output) + len(data) > RATE * 2 * 300:
                            raise VoiceError("audio: превышен лимит записи ответа 300 секунд")
                        self.output.extend(data)
                    if self.track.pcm is None:
                        if self.speaker is None:
                            import sounddevice as sd
                            self.speaker = sd.RawOutputStream(samplerate=RATE, channels=1, dtype="int16",
                                                              device=self.args.output_device)
                            self.speaker.start()
                        await asyncio.to_thread(self.speaker.write, data)
        except asyncio.CancelledError:
            raise
        except MediaStreamError:
            pass
        except Exception as error:
            self.fail(str(error) if isinstance(error, VoiceError) else "audio: ошибка воспроизведения/декодирования")

    async def stats(self):
        report = await self.pc.getStats()
        return {"packets_sent": sum(getattr(s, "packetsSent", 0) for s in report.values() if s.type == "outbound-rtp"),
                "packets_received": sum(getattr(s, "packetsReceived", 0) for s in report.values() if s.type == "inbound-rtp"),
                "decoded_samples": self.received_samples, "nonzero_samples": self.nonzero_samples,
                "transcript_chars": self.transcript_chars.copy()}

    async def close(self):
        if self.closed:
            return
        self.closed = True
        try:
            self.track.stop()
        except VoiceError:
            self.emit("microphone: ошибка закрытия, остальные ресурсы освобождаются")
        if self.ws is not None:
            try:
                await asyncio.wait_for(self.ws.send_json({"type": "session.close"}), 2)
            except Exception:
                pass
            await self.ws.close()
        for task in self.tasks:
            task.cancel()
        await asyncio.gather(*self.tasks, return_exceptions=True)
        await self.pc.close()
        if self.speaker is not None:
            self.speaker.abort()
            self.speaker.close()
            self.speaker = None
        if self.http is not None:
            await self.http.close()
        self.finished.set()


def save_wav(path, pcm):
    # Личный ответ записывается только по явному --output; не затираем файл.
    if not pcm or len(pcm) % 2:
        raise VoiceError("output: пустой или неполный PCM16")
    fd, temporary = tempfile.mkstemp(prefix=".voice-v3-", dir=path.parent)
    try:
        with os.fdopen(fd, "wb") as raw:
            with wave.open(raw, "wb") as audio:
                audio.setparams((1, 2, RATE, 0, "NONE", "not compressed"))
                audio.writeframes(pcm)
            raw.flush()
            os.fsync(raw.fileno())
        # Готовый файл публикуется атомарно; link не затирает существующий
        # target даже при гонке. mkstemp даёт 0600. Это Linux hard-link path.
        os.link(temporary, path)
    finally:
        os.unlink(temporary)  # Удаляется только собственный временный файл.


async def file_probe(args, pcm):
    session = VoiceSession(args, pcm)
    try:
        await asyncio.wait_for(session.start(), args.timeout)
        async def successful_turn():
            await session.track.file_done.wait()
            await session.assistant_done.wait()
            # Фиксированный хвост - эвристика одной реплики, не proof media drain.
            # Ранний assistant-final до конца файла не доказывает ответ на весь файл.
            await asyncio.sleep(1)
        success = asyncio.create_task(successful_turn())
        failure = asyncio.create_task(session.finished.wait())
        try:
            done, _ = await asyncio.wait([success, failure], timeout=args.max_seconds, return_when=asyncio.FIRST_COMPLETED)
            if session.failed:
                raise session.failed
            if success not in done:
                raise VoiceError("probe: недостаточно подтверждений ответа до timeout")
        finally:
            success.cancel()
            failure.cancel()
            await asyncio.gather(success, failure, return_exceptions=True)
        stats = await session.stats()
        if not (stats["packets_sent"] and stats["packets_received"] and stats["nonzero_samples"]
                and stats["transcript_chars"]["user"] and stats["transcript_chars"]["assistant"]):
            raise VoiceError("probe: нет всех подтверждений RTP + входного/выходного текста + аудио")
        save_wav(args.output, session.output)
        print(json.dumps({"transport": "WebRTC V3", "status": "success", **stats, "output": str(args.output)}))
    finally:
        await session.close()


async def console(args):
    from prompt_toolkit import PromptSession
    from prompt_toolkit.patch_stdout import patch_stdout
    session = None
    prompt = PromptSession()
    print("/voice — разговор; /voice mute — mute/unmute; /voice settings; /quit. Используйте наушники: AEC нет.")
    try:
        with patch_stdout():
            while True:
                # Во время prompt также наблюдаем сбой: микрофон закрывается сразу.
                input_task = asyncio.create_task(prompt.prompt_async("voice> "))
                monitor = asyncio.create_task(session.finished.wait()) if session else None
                try:
                    if monitor:
                        done, _ = await asyncio.wait([input_task, monitor], return_when=asyncio.FIRST_COMPLETED)
                        if monitor in done and session.failed:
                            print(session.failed)
                            await session.close()
                            session = None
                            input_task.cancel()
                            continue
                    command = (await input_task).strip()
                except (EOFError, KeyboardInterrupt):
                    break
                finally:
                    input_task.cancel()
                    if monitor:
                        monitor.cancel()
                    await asyncio.gather(input_task, *([monitor] if monitor else []), return_exceptions=True)
                if command == "/quit":
                    break
                if command == "/voice settings":
                    import sounddevice as sd
                    print(sd.query_devices())
                    print(f"V3 / CLI 0.156.1: voice={args.voice}; input={args.input_device}; output={args.output_device}; AEC=no")
                elif command == "/voice mute" and session:
                    if session.track.muted:
                        session.track.unmute()
                    else:
                        session.track.mute()
                    print("mute" if session.track.muted else "unmute")
                elif command == "/voice":
                    if session:
                        await session.close()
                        session = None
                        print("Разговор остановлен")
                    else:
                        session = VoiceSession(args)
                        try:
                            await asyncio.wait_for(session.start(), args.timeout)
                        except Exception as error:
                            print(str(error) if isinstance(error, VoiceError) else "Соединение не установлено; подробности скрыты")
                            await session.close()
                            session = None
                else:
                    print("Используйте /voice, /voice mute, /voice settings или /quit")
    finally:
        if session:
            await session.close()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--auth-file", type=Path, default=Path(os.getenv("CODEX_HOME", str(Path.home() / ".codex"))) / "auth.json")
    parser.add_argument("--voice", choices=VOICES, default="cove")
    parser.add_argument("--model", help="V3 model override; default gpt-live-1-codex как в ядре CLI 0.156.1")
    parser.add_argument("--input-device", type=int)
    parser.add_argument("--output-device", type=int)
    parser.add_argument("--file", type=Path, help="WAV 48 кГц mono PCM16 вместо микрофона")
    parser.add_argument("--output", type=Path, help="Новый WAV ответа, только для --file")
    parser.add_argument("--timeout", type=float, default=30)
    parser.add_argument("--max-seconds", type=float, default=45, help="Максимум ожидания ответа файловой пробы")
    parser.add_argument("--list-devices", action="store_true")
    parser.add_argument("--list-voices", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    for value in (args.timeout, args.max_seconds):
        if not math.isfinite(value) or not 0 < value <= 300:
            parser.error("timeout/max-seconds: конечное число от 0 до 300")
    if args.list_voices:
        print("V3, каталог CLI 0.156.1: " + ", ".join(VOICES))
        return
    if args.list_devices:
        import sounddevice as sd
        print(sd.query_devices())
        return
    if bool(args.file) != bool(args.output):
        parser.error("--file и --output задаются вместе")
    if args.output and args.output.exists():
        parser.error("--output уже существует; перезапись запрещена")
    pcm = read_wav(args.file) if args.file else None
    if args.dry_run:
        print(json.dumps({"call_url": CALL_URL, "sideband": SIDEBAND + "<call_id>", "session": call_body("<SDP>", args.voice, args.model)["session"],
                          "pcm_bytes": len(pcm) if pcm else None, "network": False, "credentials_read": False,
                          "microphone": False, "codex_process": False}, ensure_ascii=False))
        return
    asyncio.run(file_probe(args, pcm) if pcm else console(args))


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass  # asyncio.run отменяет задачи и выполняет finally/close.
    except Exception as error:
        print(str(error) if isinstance(error, VoiceError) else "V3: ошибка файла/устройства/сети; подробности скрыты", file=sys.stderr)
        raise SystemExit(1)