omnibot22/tools/markov_import.py

import logging
import asyncio
import sys
import re

from omnibot.config import ServerConfig
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()
    server_config.load("config.toml")
    # This only works on one plugin per config
    plugin_config = [
        plugin for plugin in server_config.plugins if "markov" in plugin["module"]
    ][0]
    plugin = Markov(server_config, plugin_config)

    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"]
                    plugin.add(channel, name, message)
    await plugin.save()


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