banban/talkingq-url/scripts/simulate_device_voice_exchange.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import asyncio
import json
import struct
import sys
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

import pymysql
import websockets
from pymysql.cursors import DictCursor

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config import settings


DEFAULT_WS_URL = "ws://127.0.0.1:8080/ws"


@dataclass
class DeviceContext:
    device_id: str
    serial_number: str
    child_id: int
    card_uuid: str
    inbox: asyncio.Queue[str] = field(default_factory=asyncio.Queue)
    websocket: websockets.WebSocketClientProtocol | None = None
    receiver_task: asyncio.Task | None = None


def build_runtime_args() -> argparse.Namespace:
    repo_root = Path(__file__).resolve().parents[2]
    audio_dir = repo_root / "tmp_voice_sim"

    parser = argparse.ArgumentParser(description="Simulate real websocket voice exchange between two devices")
    parser.add_argument("--ws-url", default=DEFAULT_WS_URL)
    parser.add_argument("--device-a", default="TalkingQ_device001")
    parser.add_argument("--device-b", default="TalkingQ_device002")
    parser.add_argument("--a1-audio", default=str(audio_dir / "device001_to_device002_1.mp3"))
    parser.add_argument("--b1-audio", default=str(audio_dir / "device002_to_device001_1.mp3"))
    parser.add_argument("--a2-audio", default=str(audio_dir / "device001_to_device002_2.mp3"))
    parser.add_argument("--b2-audio", default=str(audio_dir / "device002_to_device001_2.mp3"))
    return parser.parse_args()


def connect_db():
    return pymysql.connect(
        host=settings.db_host,
        port=settings.db_port,
        user=settings.db_user,
        password=settings.db_password,
        database=settings.db_name,
        charset="utf8mb4",
        cursorclass=DictCursor,
        autocommit=True,
    )


def load_device_contexts(device_ids: list[str]) -> dict[str, DeviceContext]:
    placeholders = ", ".join(["%s"] * len(device_ids))
    sql = f"""
        SELECT
          da.device_id,
          da.serial_number,
          db.child_id,
          c.card_uuid
        FROM device_auth AS da
        LEFT JOIN device_bindings AS db
          ON db.device_id = da.device_id
         AND db.status = 1
        LEFT JOIN cards AS c
          ON c.device_id = da.device_id
         AND c.status = 1
        WHERE da.device_id IN ({placeholders})
          AND da.is_active = 1
    """

    contexts: dict[str, DeviceContext] = {}
    with connect_db() as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql, device_ids)
            rows = cursor.fetchall()

    for row in rows:
        if row["child_id"] is None:
            raise RuntimeError(f"device {row['device_id']} is not bound to any child")
        if not row["card_uuid"]:
            raise RuntimeError(f"device {row['device_id']} does not have an active card")
        contexts[row["device_id"]] = DeviceContext(
            device_id=str(row["device_id"]),
            serial_number=str(row["serial_number"]),
            child_id=int(row["child_id"]),
            card_uuid=str(row["card_uuid"]),
        )

    missing = [device_id for device_id in device_ids if device_id not in contexts]
    if missing:
        raise RuntimeError(f"device context not found: {', '.join(missing)}")
    return contexts


def fetch_conversation_snapshot(child_a_id: int, child_b_id: int) -> tuple[int | None, int, int]:
    pair_key = f"{min(child_a_id, child_b_id)}:{max(child_a_id, child_b_id)}"
    with connect_db() as connection:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT id, message_count
                FROM im_conversations
                WHERE conversation_type = 1
                  AND pair_key = %s
                LIMIT 1
                """,
                (pair_key,),
            )
            row = cursor.fetchone()
            if not row:
                return None, 0, 0

            conversation_id = int(row["id"])
            message_count = int(row["message_count"])
            cursor.execute(
                "SELECT COALESCE(MAX(id), 0) AS max_message_id FROM im_messages WHERE conversation_id = %s",
                (conversation_id,),
            )
            max_row = cursor.fetchone() or {"max_message_id": 0}
            return conversation_id, message_count, int(max_row["max_message_id"] or 0)


def fetch_new_messages(conversation_id: int, min_message_id: int) -> list[dict]:
    with connect_db() as connection:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                  id,
                  conversation_id,
                  seq,
                  sender_type,
                  sender_id,
                  receiver_type,
                  receiver_id,
                  content_type,
                  media_file_key,
                  client_msg_id,
                  created_at
                FROM im_messages
                WHERE conversation_id = %s
                  AND id > %s
                ORDER BY id ASC
                """,
                (conversation_id, min_message_id),
            )
            return cursor.fetchall()


async def receiver_loop(ctx: DeviceContext) -> None:
    assert ctx.websocket is not None
    async for message in ctx.websocket:
        if isinstance(message, str):
            print(f"[{ctx.device_id}] <- {message}", flush=True)
            await ctx.inbox.put(message)
        else:
            print(f"[{ctx.device_id}] <- <binary:{len(message)}>", flush=True)


async def wait_for_text(
    ctx: DeviceContext,
    predicate: Callable[[str], bool],
    *,
    timeout: float,
    description: str,
) -> str:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    while True:
        remaining = deadline - loop.time()
        if remaining <= 0:
            raise TimeoutError(f"{ctx.device_id} timed out waiting for {description}")
        message = await asyncio.wait_for(ctx.inbox.get(), timeout=remaining)
        if predicate(message):
            return message


def create_packet(device_id: str, session_id: str, sequence_number: int, packet_type: int, audio_data: bytes) -> bytes:
    if len(session_id) != 32:
        raise ValueError("session_id must be exactly 32 ascii chars")
    return (
        device_id.encode("ascii")
        + b"\x00"
        + session_id.encode("ascii")
        + b"\x00"
        + struct.pack("<I", sequence_number)
        + struct.pack("<B", packet_type)
        + struct.pack("<I", len(audio_data))
        + audio_data
    )


async def connect_device(ctx: DeviceContext, ws_url: str) -> None:
    websocket = await websockets.connect(ws_url, ping_interval=None, max_size=None)
    ctx.websocket = websocket
    ctx.receiver_task = asyncio.create_task(receiver_loop(ctx))
    await websocket.send(json.dumps({"device_id": ctx.device_id, "serial_number": ctx.serial_number}, ensure_ascii=False))
    await wait_for_text(
        ctx,
        lambda text: json.loads(text).get("status") == "authenticated",
        timeout=10,
        description="authentication response",
    )
    print(f"[{ctx.device_id}] authenticated", flush=True)


async def send_voice_message(
    sender: DeviceContext,
    target: DeviceContext,
    audio_path: Path,
    *,
    label: str,
) -> None:
    if sender.websocket is None:
        raise RuntimeError(f"{sender.device_id} is not connected")
    if not audio_path.exists():
        raise FileNotFoundError(f"audio file not found: {audio_path}")

    audio_bytes = audio_path.read_bytes()
    session_id = uuid.uuid4().hex

    print(f"[{label}] register target card {target.card_uuid}", flush=True)
    await sender.websocket.send(f"REGISTER_TARGET_DEVICE:{target.card_uuid}")
    await wait_for_text(
        sender,
        lambda text: text.startswith("TARGET_DEVICE_REGISTERED_URL:"),
        timeout=5,
        description="target registration response",
    )

    await sender.websocket.send(create_packet(sender.device_id, session_id, 0, 1, b""))
    print(f"[{label}] start session {session_id}", flush=True)

    chunk_size = 2048
    sequence_number = 1
    for offset in range(0, len(audio_bytes), chunk_size):
        chunk = audio_bytes[offset : offset + chunk_size]
        await sender.websocket.send(create_packet(sender.device_id, session_id, sequence_number, 4, chunk))
        sequence_number += 1
        await asyncio.sleep(0.03)

    await sender.websocket.send(create_packet(sender.device_id, session_id, sequence_number, 2, b""))
    print(f"[{label}] finish session {session_id}, bytes={len(audio_bytes)}", flush=True)

    await wait_for_text(
        sender,
        lambda text: text.startswith("PROMPT_SOUND_URL:"),
        timeout=20,
        description="message stored response",
    )
    await asyncio.sleep(0.5)


async def close_device(ctx: DeviceContext) -> None:
    if ctx.websocket is not None:
        await ctx.websocket.close()
    if ctx.receiver_task is not None:
        try:
            await asyncio.wait_for(ctx.receiver_task, timeout=2)
        except Exception:
            ctx.receiver_task.cancel()


async def main() -> None:
    args = build_runtime_args()
    audio_paths = {
        "a1": Path(args.a1_audio),
        "b1": Path(args.b1_audio),
        "a2": Path(args.a2_audio),
        "b2": Path(args.b2_audio),
    }

    contexts = load_device_contexts([args.device_a, args.device_b])
    device_a = contexts[args.device_a]
    device_b = contexts[args.device_b]

    conversation_id, message_count_before, max_message_id_before = fetch_conversation_snapshot(
        device_a.child_id,
        device_b.child_id,
    )
    print(
        f"[snapshot-before] conversation_id={conversation_id} message_count={message_count_before} max_message_id={max_message_id_before}",
        flush=True,
    )

    await connect_device(device_a, args.ws_url)
    await connect_device(device_b, args.ws_url)

    try:
        await send_voice_message(device_a, device_b, audio_paths["a1"], label="A->B #1")
        await send_voice_message(device_b, device_a, audio_paths["b1"], label="B->A #1")
        await send_voice_message(device_a, device_b, audio_paths["a2"], label="A->B #2")
        await send_voice_message(device_b, device_a, audio_paths["b2"], label="B->A #2")
    finally:
        await close_device(device_a)
        await close_device(device_b)

    await asyncio.sleep(2)

    conversation_id_after, message_count_after, max_message_id_after = fetch_conversation_snapshot(
        device_a.child_id,
        device_b.child_id,
    )
    print(
        f"[snapshot-after] conversation_id={conversation_id_after} message_count={message_count_after} max_message_id={max_message_id_after}",
        flush=True,
    )
    if conversation_id_after is None:
        raise RuntimeError("child-peer conversation was not created")

    new_messages = fetch_new_messages(conversation_id_after, max_message_id_before)
    print(f"[new-messages] count={len(new_messages)}", flush=True)
    for row in new_messages:
        print(
            json.dumps(
                {
                    "id": int(row["id"]),
                    "seq": int(row["seq"]),
                    "sender_id": row["sender_id"],
                    "receiver_id": row["receiver_id"],
                    "content_type": int(row["content_type"]),
                    "media_file_key": row["media_file_key"],
                    "client_msg_id": row["client_msg_id"],
                    "created_at": row["created_at"].isoformat(sep=" "),
                },
                ensure_ascii=False,
            ),
            flush=True,
        )


if __name__ == "__main__":
    asyncio.run(main())