tilde-projects/Code/irc/madlibbot/madlibbot.py

#!/usr/bin/python
# http://wiki.shellium.org/w/Writing_an_IRC_bot_in_Python

# Import some necessary libraries.
import socket
import os
import sys
from optparse import OptionParser
import fileinput
import random
import time
import re
import operator

from .. import util
import madlib


class State:
    idle = 0  # When the bot is started and after it has finished a story
    thinking = (
        1
    )  # I intentionally speed throttle this bot so it feels a little more natural
    story_selection = 2  # When the bot is waiting for a user input for story selection
    word_input = 3  # When the bot is waiting for user input after prompted for a word


# Globals
MAX_LINE = 80  # The maximum number of characters in a line
MAX_STORIES = 5  # The maximum number of stories a user can pick from
THROTTLE_FACTOR = 1  # A factor by which all sleep event durations will be multiplied
# Allow madlbibbot to run multiple simultaneous stories
state = {}  # The current state of the bot
stories = {}  # The list of stories available to users
story = {}  # The madlib story currently being worked on
nextword = {}  # The word that the bot is currently expecting data for

parser = OptionParser()

parser.add_option(
    "-s",
    "--server",
    dest="server",
    default="127.0.0.1:6667",
    help="the server to connect to",
    metavar="SERVER",
)
parser.add_option(
    "-c",
    "--channel",
    dest="channel",
    default="#madlibs",
    help="the channel to join",
    metavar="CHANNEL",
)
parser.add_option(
    "-n",
    "--nick",
    dest="nick",
    default="madlibbot",
    help="the nick to use",
    metavar="NICK",
)

(options, args) = parser.parse_args()


def resetGlobals(channel=""):
    global state
    global stories
    global story
    global nextword

    if channel == "":
        state = {}
        stories = {}
        story = {}
        nextword = {}
    else:
        state[channel] = State.idle
        stories[channel] = []
        story[channel] = ""
        nextword[channel] = {}


def get_stories(channel, botnick):
    global state
    global stories
    state[channel] = State.thinking
    stories[channel] = madlib.find_stories(MAX_STORIES, True)
    if len(stories[channel]) == 0:
        sendmsg(channel, "Uh oh! There are no stories!")
        state[channel] = State.idle
    else:
        sendmsg(channel, "Here are a couple good stories:")
        time.sleep(1 * THROTTLE_FACTOR)
        for idx, story in enumerate(stories[channel]):
            sendmsg(channel, "[{}] {} ({} words)".format(idx, story[0], story[2]))
            time.sleep(0.5 * THROTTLE_FACTOR)
        sendmsg(
            channel,
            "Please select a story by index by saying '{}: <number>':".format(botnick),
        )
        state[channel] = State.story_selection


# Handle user input when the bot is directly addressed
def handle_bot_msg(channel, msg, botnick):
    global state
    global stories

    if channel not in state:
        state[channel] = State.idle

    saved_state = state[channel]
    state[channel] = State.thinking
    time.sleep(1 * THROTTLE_FACTOR)

    # First check if we should quit the current game
    if saved_state != State.idle and msg == "!quit":
        quit_game(channel)
    elif saved_state == State.idle and msg == "startgame":
        get_stories(channel, botnick)
    elif saved_state == State.story_selection:
        handle_story_selection(channel, msg, botnick)
    elif saved_state == State.word_input:
        handle_story_step(channel, msg)
    else:
        state[channel] = State.idle


# Handle how to quit the game
def quit_game(channel):
    resetGlobals(channel)
    sendmsg(channel, "Ok, quitting the current game.")


# Handle user input when we are in story selection mode
def handle_story_selection(channel, msg, botnick):
    global stories
    global state
    try:
        imsg = int(msg)
        if imsg < 0 or imsg > len(stories[channel]):
            sendmsg(channel, "Selection out of bounds. Try again!")
            return
        time.sleep(1 * THROTTLE_FACTOR)
        sendmsg(
            channel, "Give me a second to load up {}".format(stories[channel][imsg][0])
        )

        with open(stories[channel][imsg][1], "r") as storyFile:
            story[channel] = storyFile.read()
        stories[channel] = {}  # Clear out the saved selectable stories in memory
        story_start(channel, botnick)
    except ValueError:
        sendmsg(channel, "Invalid selection. Try again!")
        state[channel] = State.story_selection


# Handle when a story is being started
def story_start(channel, botnick):
    global story
    global state
    global nextword

    state[channel] = State.thinking
    sendmsg(
        channel,
        "Alright! Let's get started! Say '{}: <word>' to give me words.".format(
            botnick
        ),
    )
    nextword[channel] = madlib.find_next_word(story[channel], True)
    time.sleep(0.5 * THROTTLE_FACTOR)
    sendmsg(channel, "Give me {}:".format(nextword[channel][1]))
    state[channel] = State.word_input


# Handle user input when we have asked the user for input and are expecting a
# response
def handle_story_step(channel, msg):
    global state
    global story
    global nextword

    state[channel] = State.thinking
    word = nextword[channel]  # madlib.find_next_word(story[channel])
    if word is not None:
        story[channel] = madlib.replace_word(story[channel], nextword[channel][0], msg)
    time.sleep(1 * THROTTLE_FACTOR)

    nextword[channel] = madlib.find_next_word(story[channel], True)
    if nextword[channel] is None:
        finish_story(channel)
        return
    # else
    count = madlib.count_words(story[channel])
    sendmsg(
        channel,
        "Thanks! Now give me {} ({} words left)".format(nextword[channel][1], count),
    )
    state[channel] = State.word_input


# Finish the story
def finish_story(channel):
    global state
    global story

    sendmsg(channel, "Ok, here's the story...")
    sendmsg(channel, "=" * MAX_LINE)
    for line in story[channel].splitlines():
        for part in madlib.yield_lines(line, MAX_LINE):
            time.sleep(0.6 * THROTTLE_FACTOR)
            sendmsg(channel, part)
    padlen = (MAX_LINE - 9) / 2
    mod = (MAX_LINE - 9) % 2
    sendmsg(channel, "=" * padlen + " THE END " + "=" * (padlen + mod))

    story[channel] = ""
    state[channel] = State.idle


# System things
def joinchan(chan):
    global state
    state[chan] = State.idle
    ircsock.send("JOIN {}\n".format(chan))


def rollcall(channel, botnick):
    global state
    if channel not in state:
        state[channel] = State.idle

    if state[channel] == State.idle:
        sendmsg(
            channel,
            "Do you like MadLibs? Start a collaborative story by saying '{}: startgame'".format(
                botnick
            ),
        )
    else:
        sendmsg(
            channel,
            "A game is already in progress. Say '{}: <word>' to provide me with the next word or '{}: !quit' to stop the current game".format(
                botnick, botnick
            ),
        )


def listen(botnick):
    botmsgre = re.compile(
        "^{}\:?\s*(.*)$".format(botnick)
    )  # re to strip the bot's name from a message
    while 1:
        ircmsg = ircsock.recv(2048)
        ircmsg = ircmsg.strip("\n\r")

        if ircmsg[:4] == "PING":
            ping(ircmsg.split(" ")[1])

        formatted = util.format_message(ircmsg)
        if "" == formatted:
            continue
        # print formatted

        split = formatted.split("\t")
        msgtime = split[0]
        user = split[1]
        command = split[2]
        channel = split[3]
        message = split[4]

        if (
            message.startswith("!rollcall") == True
            or message.startswith("!help") == True
        ):
            rollcall(channel, botnick)
        elif message.startswith(botnick) == True:
            botmsg = botmsgre.match(message).group(1)
            handle_bot_msg(channel, botmsg, botnick)

        sys.stdout.flush()
        time.sleep(1 * THROTTLE_FACTOR)


ircsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
util.connect(ircsock, options)
listen(options.nick)