"""文字おこしくん Subscribe プロトコル v1 の参照実装（Python 3.10+、websockets 13 以降）。

    pip install websockets
    python mojiokoshi_subscribe.py wss://<host>/v1/subscribe <transcript:read のトークン>

確定した発話（final / refined）を 1 行ずつ出す。守っている決まり: トークンは hello で送る /
welcome で状態を丸ごと置き換える / 発話は utteranceId ごとに revision が最大のものだけ残す（追記しない）/
ping の定期送信 / 指数バックオフ + ジッターの再接続 / 再接続しても直らないエラーと session_ended で止まる。
"""

import asyncio
import json
import random
import sys

import websockets

STOP_CODES = {"unauthorized", "protocol_version_unsupported", "session_ended", "superseded", "quota_exceeded"}


class Stop(Exception):
    pass


async def run(url: str, token: str) -> None:
    speakers: dict[str, str] = {}  # speakerId -> displayName
    utterances: dict[str, dict] = {}  # utteranceId -> 最新の改訂
    printed: dict[str, int] = {}  # utteranceId -> 出力済みの revision
    attempt = 0

    def show(u: dict) -> None:
        # refined は final を置き換える改訂。同じ発話を 2 回出すときは「(修正)」を付ける。
        if u["stage"] == "partial" or printed.get(u["utteranceId"], -1) >= u["revision"]:
            return
        who = f"track:{u['trackId']}" if u["speaker"]["kind"] == "unattributed" else speakers.get(u["speaker"]["speakerId"], "?")
        mark = " (修正)" if u["utteranceId"] in printed else ""
        printed[u["utteranceId"]] = u["revision"]
        print(f"{who}: {u['text']}{mark}", flush=True)

    while True:
        loop = asyncio.get_running_loop()
        opened_at = None
        try:
            # ブラウザと同じく JSON の ping を使う（ライブラリの ping フレームは切っておく）。
            async with websockets.connect(url, ping_interval=None) as ws:
                await ws.send(json.dumps({"type": "hello", "protocolVersion": 1, "token": token}))
                pinger = asyncio.create_task(ping_loop(ws))
                try:
                    async for raw in ws:
                        if not isinstance(raw, str):
                            continue
                        m = json.loads(raw)
                        t = m.get("type")
                        if t == "welcome":
                            opened_at = loop.time()
                            speakers = {e["speaker"]["speakerId"]: e["speaker"]["displayName"] for e in m["speakers"]}
                            utterances = {}
                            for u in m["recent"]:
                                upsert(utterances, u)
                        elif t == "speaker_upserted":
                            speakers[m["speaker"]["speakerId"]] = m["speaker"]["displayName"]
                        elif t == "utterance":
                            if upsert(utterances, m["utterance"]):
                                show(m["utterance"])
                        elif t == "session_ended":
                            raise Stop("session ended")
                        elif t == "error":
                            e = m["error"]
                            if e["fatal"] and (e["code"] in STOP_CODES or e["code"] == "invalid_message"):
                                raise Stop(f"{e['code']}: {e['message']}")
                            if e["fatal"] and e["code"] == "rate_limited":
                                attempt = 6  # 次の待ち時間を上限にする
                finally:
                    pinger.cancel()
        except Stop as stop:
            print(f"停止: {stop}", file=sys.stderr)
            return
        except (OSError, websockets.WebSocketException) as error:
            print(f"切断: {error!r}", file=sys.stderr)
        # 10 秒以上つながっていたら待ち時間を最初に戻す。0.5 秒から倍々、上限 30 秒、ジッター 50%。
        if opened_at is not None and loop.time() - opened_at > 10:
            attempt = 0
        delay = min(30.0, 0.5 * 2**attempt) * (1 - 0.5 * random.random())
        attempt += 1
        await asyncio.sleep(delay)


def upsert(utterances: dict[str, dict], u: dict) -> bool:
    """utteranceId ごとに revision 最大の 1 件だけ残す。変化したら True。"""
    known = utterances.get(u["utteranceId"])
    if known is not None and known["revision"] >= u["revision"]:
        return False
    utterances[u["utteranceId"]] = u
    return True


async def ping_loop(ws) -> None:
    while True:
        await asyncio.sleep(15)
        await ws.send('{"type":"ping"}')


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("使い方: python mojiokoshi_subscribe.py <wss://…/v1/subscribe> <token>", file=sys.stderr)
        sys.exit(2)
    asyncio.run(run(sys.argv[1], sys.argv[2]))