bililive-touhou-tts/server.py

"""Bilibili Live → Yukkuri TTS — Web UI Server."""

import argparse
import asyncio
import json
import logging
import os
import sys
import time
from pathlib import Path

from aiohttp import web

sys.path.insert(0, os.path.dirname(__file__))

from danmaku_handler import DanmakuClient
from chinese2kana import chinese_to_kana, filter_kana
from tts import TTSBridge
from audio_player import play_wav_async, get_output_devices
from bili_login import BiliLoginSession
from event_rules import DEFAULT_RULES, load_rules, save_rules, rules_path

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

def _get_static_dir() -> Path:
    if getattr(sys, "frozen", False):
        return Path(sys._MEIPASS) / "static"
    return Path(__file__).parent / "static"

STATIC_DIR = _get_static_dir()
event_rules = load_rules()


class TTSService:
    def __init__(self):
        self._bridge: TTSBridge | None = None
        self._danmaku_client: DanmakuClient | None = None
        self._queue: asyncio.Queue | None = None
        self._shutdown_event: asyncio.Event | None = None
        self._worker_task: asyncio.Task | None = None
        self._cookies: dict[str, str] | None = None
        self._lifecycle_lock = asyncio.Lock()
        self.running = False
        self.start_time: float = 0
        self.messages_processed = 0
        self.recent_messages: list[str] = []
        self.current_config: dict = {}
        self._volume: int = 100

    def set_cookies(self, cookies: dict[str, str] | None) -> None:
        self._cookies = cookies

    async def start(self, config: dict) -> None:
        async with self._lifecycle_lock:
            if self.running:
                raise RuntimeError("Already running")
            room_id = config["room_id"]
            voice = config.get("voice", "f1")
            speed = config.get("speed", 100)
            self._volume = config.get("volume", 100)
            convert_numbers = config.get("convert_numbers", True)
            self.current_config = config
            self.recent_messages = []
            self.messages_processed = 0
            self.start_time = time.time()
            self._queue = asyncio.Queue(maxsize=256)
            self._shutdown_event = asyncio.Event()
            self._bridge = TTSBridge(voice=voice, speed=speed)
            await self._bridge.start()
            self._danmaku_client = DanmakuClient(
                room_id=room_id, queue=self._queue,
                cookies=self._cookies, rules=event_rules,
            )
            await self._danmaku_client.start()
            self.running = True
            self._worker_task = asyncio.create_task(self._tts_worker(convert_numbers))

    async def stop(self) -> None:
        async with self._lifecycle_lock:
            if not self.running:
                return
            self.running = False
            if self._shutdown_event:
                self._shutdown_event.set()
            if self._worker_task:
                self._worker_task.cancel()
                try:
                    await self._worker_task
                except asyncio.CancelledError:
                    pass
            if self._danmaku_client:
                await self._danmaku_client.stop()
            if self._bridge:
                await self._bridge.stop()
            self._queue = None
            self._bridge = None
            self._danmaku_client = None

    async def _tts_worker(self, convert_numbers: bool) -> None:
        total_in = 0
        total_out = 0
        while not self._shutdown_event.is_set():
            try:
                text = await asyncio.wait_for(self._queue.get(), timeout=0.5)
            except asyncio.TimeoutError:
                continue
            total_in += 1
            try:
                kana = chinese_to_kana(text, convert_numbers=convert_numbers)
                kana = filter_kana(kana)
                if not kana.strip():
                    self._add_message(f"SKIP(empty kana): {text}")
                    self._queue.task_done()
                    continue
                self._add_message(f"[{total_in}] {text}")
                logger.info("TTS #%d: %s -> %s", total_in, text[:40], kana[:40])
                wav_data = await self._bridge.synthesize(kana)
                await play_wav_async(wav_data, self._volume)
                total_out += 1
                self.messages_processed += 1
            except asyncio.TimeoutError:
                self._add_message(f"TIMEOUT: {text[:60]}")
                logger.error("TTS timeout for: %s", text[:60])
            except Exception as e:
                self._add_message(f"ERR: {text[:60]} | {e}")
                logger.error("TTS error: %s", str(e)[:200])
            finally:
                self._queue.task_done()

    def _add_message(self, msg: str) -> None:
        self.recent_messages.append(msg)
        if len(self.recent_messages) > 200:
            self.recent_messages = self.recent_messages[-200:]

    def status(self) -> dict:
        uptime = time.time() - self.start_time if self.start_time else 0
        return {
            "running": self.running,
            "uptime": int(uptime),
            "messages_processed": self.messages_processed,
            "recent_messages": self.recent_messages[-50:],
            "config": self.current_config,
        }


tts_service = TTSService()
login_session = BiliLoginSession()


async def index_handler(request: web.Request) -> web.Response:
    return web.FileResponse(STATIC_DIR / "index.html")


async def api_start(request: web.Request) -> web.Response:
    try:
        config = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)
    room_id = config.get("room_id")
    if not room_id:
        return web.json_response({"error": "room_id is required"}, status=400)
    try:
        room_id = int(room_id)
    except (ValueError, TypeError):
        return web.json_response({"error": "room_id must be an integer"}, status=400)
    config["room_id"] = room_id
    config.setdefault("voice", "f1")
    config.setdefault("speed", 100)
    config.setdefault("volume", 100)
    config.setdefault("convert_numbers", True)
    if tts_service.running:
        try:
            await tts_service.stop()
        except Exception:
            pass
    tts_service.set_cookies(login_session.cookies if login_session.is_logged_in else None)
    try:
        await tts_service.start(config)
    except Exception as e:
        logger.exception("Failed to start TTS service")
        return web.json_response({"error": str(e)}, status=500)
    return web.json_response({"status": "started"})


async def api_stop(request: web.Request) -> web.Response:
    try:
        await tts_service.stop()
    except Exception as e:
        logger.exception("Failed to stop TTS service")
        return web.json_response({"error": str(e)}, status=500)
    return web.json_response({"status": "stopped"})


async def api_status(request: web.Request) -> web.Response:
    s = tts_service.status()
    s["login"] = {
        "logged_in": login_session.is_logged_in,
        "username": login_session.username,
        "uid": login_session.uid,
    }
    return web.json_response(s)


async def api_devices(request: web.Request) -> web.Response:
    try:
        devices = get_output_devices()
        return web.json_response({"devices": devices})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


async def api_qr_generate(request: web.Request) -> web.Response:
    try:
        result = await login_session.generate_qrcode()
        return web.json_response(result)
    except Exception as e:
        logger.exception("QR generate failed")
        return web.json_response({"error": str(e)}, status=500)


async def api_qr_poll(request: web.Request) -> web.Response:
    try:
        result = await login_session.poll_login()
        return web.json_response(result)
    except Exception as e:
        logger.exception("QR poll failed")
        return web.json_response({"error": str(e)}, status=500)


async def api_login_status(request: web.Request) -> web.Response:
    return web.json_response({
        "logged_in": login_session.is_logged_in,
        "username": login_session.username,
        "uid": login_session.uid,
    })


async def api_logout(request: web.Request) -> web.Response:
    login_session.clear()
    return web.json_response({"status": "logged_out"})


async def api_rules_get(request: web.Request) -> web.Response:
    return web.json_response(event_rules)


async def api_rules_post(request: web.Request) -> web.Response:
    try:
        new_rules = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    from event_rules import ALL_EVENTS, DEFAULT_RULES
    merged = {
        "enabled": dict(DEFAULT_RULES["enabled"]),
        "templates": dict(DEFAULT_RULES["templates"]),
        "skip_free_gift": True,
    }
    if isinstance(new_rules.get("enabled"), dict):
        for k in ALL_EVENTS:
            if isinstance(new_rules["enabled"].get(k), bool):
                merged["enabled"][k] = new_rules["enabled"][k]
    if isinstance(new_rules.get("templates"), dict):
        for k in ALL_EVENTS:
            if isinstance(new_rules["templates"].get(k), str):
                merged["templates"][k] = new_rules["templates"][k]
    if isinstance(new_rules.get("skip_free_gift"), bool):
        merged["skip_free_gift"] = new_rules["skip_free_gift"]

    try:
        save_rules(merged)
    except OSError as e:
        return web.json_response({"error": f"Failed to save rules: {e}"}, status=500)

    # Mutate in place so running handlers (which hold the same dict reference)
    # pick up the new rules immediately without restarting TTS.
    event_rules.clear()
    event_rules.update(merged)
    return web.json_response({"status": "saved", "rules": event_rules})


def create_app() -> web.Application:
    app = web.Application()
    app.router.add_get("/", index_handler)
    app.router.add_post("/api/start", api_start)
    app.router.add_post("/api/stop", api_stop)
    app.router.add_get("/api/status", api_status)
    app.router.add_get("/api/devices", api_devices)
    app.router.add_get("/api/qr/generate", api_qr_generate)
    app.router.add_get("/api/qr/poll", api_qr_poll)
    app.router.add_get("/api/login/status", api_login_status)
    app.router.add_get("/api/logout", api_logout)
    app.router.add_get("/api/rules", api_rules_get)
    app.router.add_post("/api/rules", api_rules_post)
    return app


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Bilibili Live → Yukkuri TTS Web UI")
    parser.add_argument("--port", "-p", type=int, default=8080)
    parser.add_argument("--host", type=str, default="127.0.0.1")
    parser.add_argument("--debug", "-d", action="store_true")
    return parser.parse_args()


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",
    )
    logging.getLogger("aiohttp.access").setLevel(logging.WARNING)

    app = create_app()
    print(f"\n  Bilibili Live → Yukkuri TTS Web UI")
    print(f"  Open: http://{args.host}:{args.port}\n")
    web.run_app(app, host=args.host, port=args.port, print=lambda *a: None)
    return 0


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