import aoclib

DAY = 19
TEST_SOLUTION_PART1 = 3
TEST_SOLUTION_PART2 = 12
rule_dict = {}


def populate_rule_dict(raw_input):
    for line in raw_input:
        rule_id, rule = line.split(": ")
        if "|" in rule:
            rule_dict[rule_id] = rule.split(" | ")
        else:
            rule_dict[rule_id] = rule


def join_messages(msglist):
    if len(msglist) == 1:
        return msglist[0]

    if len(msglist) > 2:
        msglist = [msglist[0], join_messages(msglist[1:])]

    answerlist = []
    for x in msglist[0]:
        for y in msglist[1]:
            answerlist.append("%s%s" % (x, y))

    return answerlist


def get_valid_messages(rule_id="0"):
    rule = rule_dict[rule_id]
    if '"' in rule:
        return [rule[1]]

    if not isinstance(rule, list):
        submessages = [get_valid_messages(subrule) for subrule in rule.split(" ")]
        submessages = join_messages(submessages)
    else:
        submessages = []
        for choice in rule:
            choicemessages = [get_valid_messages(subrule) for subrule in choice.split(" ")]
            submessages.extend(join_messages(choicemessages))

    return submessages


def part1(test_mode=False):
    my_input = aoclib.getMultiLineInputAsArray(day=DAY, test=test_mode)

    populate_rule_dict(my_input[0])
    # split rule 0 to save time
    # rule 0 == 8 11
    valid_messages_8 = get_valid_messages("8")
    valid_messages_11 = get_valid_messages("11")
    len_8 = len(valid_messages_8[0])
    len_11 = len(valid_messages_11[0])
    valid_messages_8 = set(valid_messages_8)
    valid_messages_11 = set(valid_messages_11)

    counter = 0
    for message in my_input[1]:
        if message[:len_8] in valid_messages_8:
            message = message[len_8:]

        if message[:len_11] in valid_messages_11:
            message = message[len_11:]

        if len(message) == 0:
            counter += 1

    return counter


def part2(test_mode=False):
    my_input = aoclib.getMultiLineInputAsArray(day=DAY, test=test_mode)

    populate_rule_dict(my_input[0])
    valid_messages_42 = get_valid_messages("42")
    valid_messages_31 = get_valid_messages("31")

    len_8 = len(valid_messages_42[0])
    len_31 = len(valid_messages_31[0])
    valid_messages_42 = set(valid_messages_42)
    valid_messages_31 = set(valid_messages_31)

    # rule 0 == 8 11
    # rule 8 == 42 | 42 8 (so just a multiple of 42)
    # rule 11 == 42 31 | 42 11 31 (so any multiple of 42 followed the same multiple of 31)
    # so a valid message consists of a multiple of 42 followed by a multiple of 42 followed by a multiple of 31
    # just make sure you find at least one more 42 than 31 (to accomodate for the 8) and nothing else
    counter = 0
    for message in my_input[1]:
        count42 = 0
        while message[:len_8] in valid_messages_42:
            message = message[len_8:]
            count42 += 1

        if len(message) == 0:
            continue

        count31 = 0
        while message[:len_31] in valid_messages_31:
            message = message[len_31:]
            count31 += 1

        if len(message) == 0 and count42 > count31:
            counter += 1

    return counter