bililive-touhou-tts/main.py

"""bililive-touhou-tts — Bilibili Live Danmaku → Yukkuri TTS Reader.

Usage:
    python main.py --room-id 12345
    python main.py --room-id 12345 --voice f2 --speed 120
"""

import argparse
import asyncio
import logging
import signal
import sys
from pathlib import Path

from danmaku_handler import DanmakuClient
from chinese2kana import chinese_to_kana
from tts import TTSBridge
from audio_player import play_wav_async, get_output_devices

logger = logging.getLogger("bililive-tts")

DEFAULT_VOICE = "f1"
DEFAULT_SPEED = 100


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Bilibili Live Danmaku → Yukkuri TTS Reader"
    )
    parser.add_argument(
        "--room-id", "-r",
        type=int,
        required=True,
        help="Bilibili live room ID (from the URL)",
    )
    parser.add_argument(
        "--voice", "-v",
        type=str,
        default=DEFAULT_VOICE,
        choices=["f1", "f2", "m1", "m2", "dvd", "imd1", "jgr", "r1"],
        help=f"Yukkuri voice type (default: {DEFAULT_VOICE})",
    )
    parser.add_argument(
        "--speed", "-s",
        type=int,
        default=DEFAULT_SPEED,
        help=f"Speech speed, 50-300 (default: {DEFAULT_SPEED})",
    )
    parser.add_argument(
        "--sessdata",
        type=str,
        default="",
        help="Bilibili SESSDATA cookie for authenticated access (optional)",
    )
    parser.add_argument(
        "--format", "-f",
        type=str,
        default="{uname}: {msg}",
        help="Danmaku message format. Variables: {uname}, {msg}, {uid}",
    )
    parser.add_argument(
        "--no-numbers",
        action="store_true",
        help="Do not convert Arabic numerals to Chinese words",
    )
    parser.add_argument(
        "--list-devices",
        action="store_true",
        help="List available audio output devices and exit",
    )
    parser.add_argument(
        "--debug", "-d",
        action="store_true",
        help="Enable debug logging",
    )
    return parser.parse_args()


async def tts_worker(
    queue: asyncio.Queue,
    bridge: TTSBridge,
    convert_numbers: bool,
    shutdown_event: asyncio.Event,
) -> None:
    """Consume danmaku messages from the queue, convert to kana, and play TTS."""
    logger.info("TTS worker started.")

    while not shutdown_event.is_set():
        try:
            text = await asyncio.wait_for(queue.get(), timeout=1.0)
        except asyncio.TimeoutError:
            continue

        try:
            kana = chinese_to_kana(text, convert_numbers=convert_numbers)
            logger.info("Speaking: %s → %s", text, kana)

            wav_data = await bridge.synthesize(kana)
            await play_wav_async(wav_data)

        except Exception:
            logger.exception("Error processing danmaku: %s", text)

        finally:
            queue.task_done()

    logger.info("TTS worker stopped.")


async def main_async(args: argparse.Namespace) -> int:
    if args.list_devices:
        devices = get_output_devices()
        print("Available audio output devices:")
        for dev in devices:
            print(f"  [{dev['index']}] {dev['name']} ({dev['channels']}ch)")
        return 0

    if not (50 <= args.speed <= 300):
        logger.error("Speed must be between 50 and 300.")
        return 1

    queue: asyncio.Queue = asyncio.Queue(maxsize=256)
    shutdown_event = asyncio.Event()

    bridge = TTSBridge(voice=args.voice, speed=args.speed)

    danmaku_client = DanmakuClient(
        room_id=args.room_id,
        queue=queue,
        sessdata=args.sessdata,
        message_format=args.format,
    )

    async def handle_signal() -> None:
        """Wait for shutdown signal."""
        loop = asyncio.get_running_loop()
        sig_event = asyncio.Event()

        def _handler():
            sig_event.set()

        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.add_signal_handler(sig, _handler)
            except NotImplementedError:
                pass

        await sig_event.wait()
        logger.info("Shutdown signal received.")
        shutdown_event.set()

    try:
        await bridge.start()
        await danmaku_client.start()

        worker_task = asyncio.create_task(
            tts_worker(queue, bridge, not args.no_numbers, shutdown_event)
        )

        await handle_signal()

        worker_task.cancel()
        try:
            await worker_task
        except asyncio.CancelledError:
            pass

    except Exception:
        logger.exception("Fatal error")
        return 1
    finally:
        await danmaku_client.stop()
        await bridge.stop()

    return 0


def main() -> int:
    args = parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%H:%M:%S",
    )

    try:
        return asyncio.run(main_async(args))
    except KeyboardInterrupt:
        return 0


if __name__ == "__main__":
    sys.exit(main())