omnibot22/tools/markov_import.py

import logging
import asyncio
import sys
import re

from omnibot.config import ServerConfig
from omnibot.bot import Bot
from plugins.markov import Markov


async def main():
    """
    Hacky "load my IRC logs" script
    """

    # TODO - add config path selection
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s - %(name)-12s - %(levelname)-8s - %(message)s",
    )
    log = logging.getLogger(__name__)

    channel = sys.argv[1]
    files = sys.argv[2:]
    server_config = ServerConfig()
    with open("config.toml") as fp:
        server_config.load(fp)
    # This only works on one plugin per config
    bot = Bot(server_config)
    plugin = [plugin for plugin in bot.plugins if isinstance(plugin, Markov)][0]

    LINE_RE = re.compile(r"^\[[^\]]+\] <(?P<name>[^>]+)> (?P<message>.+)$")

    for fname in files:
        log.info("Loading %s", fname)
        with open(fname) as fp:
            lines = list(fp)
            for line in lines:
                if mat := LINE_RE.search(line):
                    name = mat["name"]
                    message = mat["message"].strip()
                    if name != server_config.nick and message and message[0] != "!":
                        plugin.add(channel, name, message, commit=False)
    await plugin.save()


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