vibe-bot/vibe_bot/app.py

"""Composition root: owns every singleton, the four services, and the bot."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING

import discord
from discord import Message
from discord.ext import commands

from vibe_bot import llm_client
from vibe_bot.config import TTS_MODEL_PATH, TTS_VOICES_PATH
from vibe_bot.database import ChatDatabase, CustomBotManager
from vibe_bot.llm_client import ToolRegistry
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import ConversationService
from vibe_bot.services.image_service import ImageService
from vibe_bot.services.speech_service import SpeechService
from vibe_bot.tts import TTSEngine

if TYPE_CHECKING:
    from discord.ext.commands import Bot
    from discord.ext.commands import Context as CommandsContext

logger = logging.getLogger(__name__)


def configure_logging() -> None:
    """Configure root logging (the single ``basicConfig`` in the codebase)."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )


@dataclass
class App:
    """The composition root: every singleton and the four services."""

    db: ChatDatabase
    manager: CustomBotManager
    registry: ToolRegistry
    tts: TTSEngine | None
    chat: ChatService
    image: ImageService
    speech: SpeechService
    conversation: ConversationService
    bot_cache: dict[str, tuple[str, str]]


def _build_bot_cache(manager: CustomBotManager) -> dict[str, tuple[str, str]]:
    """Snapshot the custom bots as name -> (system_prompt, creator)."""
    return {
        name: (system_prompt, creator)
        for name, system_prompt, creator in manager.list_custom_bots()
    }


def invalidate_bot_cache(app: App) -> None:
    """Rebuild ``app.bot_cache`` from the database after a change."""
    app.bot_cache = _build_bot_cache(app.manager)


def create_app() -> App:
    """Build the App: singletons, tolerant TTS init, and the four services."""
    db = ChatDatabase()
    manager = CustomBotManager()
    registry = llm_client.get_tool_registry()

    engine: TTSEngine | None = None
    try:
        engine = TTSEngine(TTS_MODEL_PATH, TTS_VOICES_PATH)
        logger.info("TTS engine initialized successfully")
    except Exception:
        logger.exception("Failed to initialize TTS engine")
        logger.info(
            "Make sure kokoro-v1.0.onnx and voices-v1.0.bin are in the project directory",
        )

    return App(
        db=db,
        manager=manager,
        registry=registry,
        tts=engine,
        chat=ChatService(db, registry),
        image=ImageService(db, discord.File),
        speech=SpeechService(db, manager, engine, discord.File),
        conversation=ConversationService(manager),
        bot_cache=_build_bot_cache(manager),
    )


# Module-level holders wired by build_bot(); the event handlers below are
# module-level so they stay importable and testable.
_app: App | None = None
_bot: commands.Bot | None = None


def _require_app() -> App:
    """The App wired by build_bot(); handlers must not run before it."""
    if _app is None:
        msg = "App is not initialized; build_bot(app) must be called first."
        raise RuntimeError(msg)
    return _app


def _require_bot() -> commands.Bot:
    """The Bot wired by build_bot(); event handlers must not run before it."""
    if _bot is None:
        msg = "Bot is not initialized; build_bot(app) must be called first."
        raise RuntimeError(msg)
    return _bot


async def on_ready() -> None:
    """Log when the bot is ready and logged in."""
    bot = _require_bot()
    logger.info("Bot is starting up...")
    logger.info("Bot logged in as %s", bot.user)


async def on_message(message: Message) -> None:
    """Handle incoming messages for custom bot command detection."""
    app = _require_app()
    bot = _require_bot()

    if message.author == bot.user:
        return

    if not message.content.startswith("!"):
        return

    message_content = message.content.lower()
    logger.debug(
        "Processing message from user %s (chars=%d)",
        message.author.id,
        len(message_content),
    )

    for bot_name, (system_prompt, _creator) in app.bot_cache.items():
        if message_content.startswith(f"!{bot_name} "):
            logger.info(
                "Custom bot %r triggered by user %s", bot_name, message.author.id
            )

            user_message = message.content[len(f"!{bot_name} ") :]
            logger.debug(
                "Extracted user message for bot %r (chars=%d)",
                bot_name,
                len(user_message),
            )

            response_prefix = f"{bot_name} response"
            logger.info("Sending request to chat service for bot %r", bot_name)
            ctx = await bot.get_context(message)
            await app.chat.handle(
                ctx=ctx,
                bot_name=bot_name,
                message=user_message,
                system_prompt=system_prompt,
                response_prefix=response_prefix,
            )
            return

    # If no custom bot matched, call the default event handler
    await bot.process_commands(message)


async def on_command_error(
    ctx: CommandsContext[Bot],
    error: commands.CommandError,
) -> None:
    """Send a friendly message when a command hits its per-user cooldown."""
    if isinstance(error, commands.CommandOnCooldown):
        retry_after = max(1, round(error.retry_after))
        await ctx.send(f"You're using that too quickly, try again in {retry_after}s.")
        return
    logger.exception("Unhandled command error in %s: %s", ctx.command, error)


def build_bot(app: App) -> commands.Bot:
    """Create the Discord bot with all event and command handlers attached."""
    global _app, _bot
    _app = app

    intents = discord.Intents.default()
    intents.message_content = True
    intents.members = True
    intents.presences = True
    bot = commands.Bot(command_prefix="!", intents=intents)
    _bot = bot

    bot.event(on_ready)
    bot.event(on_message)
    bot.event(on_command_error)

    # Imported here (not at module top) to break the import cycle.
    from vibe_bot.commands import register_all

    register_all(bot, app)
    return bot