#! /usr/bin/env python3
from itertools import product
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple

MAX_DEPTH = 10


class Rule(object):
    _rule_body: str
    key: int
    _values: List[str]
    _depth: int

    def __init__(self, key: int, rule_body: str):
        self.key = key
        self._rule_body = rule_body
        self._values = []
        self._depth = 0

    def validate(self, s: str) -> bool:
        return s in self.values()

    def values(self) -> List[str]:
        if not self._values:
            if self._depth >= MAX_DEPTH:
                # Break out of recursion
                return []

            self._depth += 1

            values: List[str] = []
            for sub_rule in self._rule_body.split(" | "):
                rule_combo: "Rule" = ValueRule(-1, [""])
                for rule_key in sub_rule.split(" "):
                    rule_combo += RULE_INDEX[int(rule_key)]

                values += rule_combo.values()

            self._values = values

        return self._values

    def __add__(self, other) -> "Rule":
        values = [
            "".join(v)
            for v in product(self.values(), other.values())
        ]
        return ValueRule(
            -1,
            values,
        )

    def __mul__(self, other) -> List[str]:
        raise NotImplementedError()

    def __repr__(self) -> str:
        return f"<Rule key={self.key} rule={self._rule_body}>"


class ValueRule(Rule):
    key: int

    def __init__(self, key: int, values: List[str]):
        self.key = key
        self._values = values

    def values(self) -> List[str]:
        return self._values

    def __repr__(self) -> str:
        return f"<ValueRule key={self.key} values={self.values()}>"


def parse_rule(line) -> Tuple[int, Rule]:
    rule_index = line.find(":")
    if rule_index <= 0:
        raise ValueError(f"Could not find rule number in {line}")
    rule_key = int(line[:rule_index])
    rule_body = line[rule_index+2:]
    rule: Optional[Rule] = None
    if rule_body[0] == '"':
        rule = ValueRule(rule_key, [rule_body[1:-1]])
    else:
        rule = Rule(rule_key, rule_body)

    return rule_key, rule


RULE_INDEX: Dict[int, Rule] = {}


def test():
    test_input = (
        "0: 4 1 5",
        "1: 2 3 | 3 2",
        "2: 4 4 | 5 5",
        "3: 4 5 | 5 4",
        '4: "a"',
        '5: "b"',
    )
    RULE_INDEX.update({
        parse_rule(line)
        for line in test_input
    })

    print(RULE_INDEX)

    for rule in RULE_INDEX.values():
        print(rule.key, rule.values())

    test_eval = (
        "ababbb",
        "bababa",
        "abbbab",
        "aaabbb",
        "aaaabbb",
    )

    rule_0 = RULE_INDEX[0]
    matched = 0

    for t in test_eval:
        if rule_0.validate(t):
            matched += 1

    print(f"Total matched is {matched}")


def part1():
    rules_parsed = False
    matched = 0

    with open("input.txt") as f:
        for line in f:
            line = line.strip()
            print(line, end="")
            if line == "":
                rules_parsed = True
                print("DONE PARSING RULES")
            elif not rules_parsed:
                key, rule = parse_rule(line)
                RULE_INDEX[key] = rule
                print(" Parsed")
            elif RULE_INDEX[0].validate(line):
                matched += 1
                print(" OK")
            else:
                print(" Not OK")

    print(f"Total matched is {matched}")


def chunk(a, n):
    return (a[i:i+n] for i in range(0, len(a), n))


def count_prefix(value: str, prefix: str) -> int:
    count = 0
    for c in chunk(value, len(prefix)):
        if c == prefix:
            count += 1
        else:
            break

    return count


class Rule0(Rule):
    def __init__(self):
        super().__init__(0, "")

    def validate(self, value: str) -> bool:
        r42 = RULE_INDEX[42]
        r31 = RULE_INDEX[31]

        # print(42, "\n".join(r42.values()))
        # print(31, "\n".join(r31.values()))

        size = None
        for v in r42.values():
            if size is None:
                size = len(v)
            assert len(v) == size
        for v in r31.values():
            assert len(v) == size

        match42, match31 = 0, 0
        for c in chunk(value, size):
            if match31 == 0 and r42.validate(c):
                match42 += 1
            elif match42 != 0 and r31.validate(c):
                match31 += 1
            else:
                return False

        # if value == "abbbbbabbbaaaababbaabbbbabababbbabbbbbbabaaaa":
        #     from pdb import set_trace; set_trace()

        if match42 > match31 and match31 > 0:
            return True

        return False

        # if RULE_INDEX[0].validate(value):
        #     from pdb import set_trace; set_trace()

        # Check beginning starts with "8"
        for prefix1 in r42.values():
            prefix_count = count_prefix(value, prefix1)

            if prefix_count == 0:
                # Not even one chunk was matched, this is a bust
                continue

            # Trim first prefix out
            back_half = value[len(prefix1) * prefix_count:]

            # Check following "11" match
            for prefix2, suffix in product(r42.values(), r31.values()):
                remain = back_half

                if RULE_INDEX[0].validate(value) and prefix1 == "bbabb" and prefix2 == "bbaab" and suffix == "aabba":
                    from pdb import set_trace; set_trace()

                # If prefixes aren't the same, look for the new one
                if prefix1 != prefix2:
                    prefix_count = count_prefix(remain, prefix2)
                    # Trim the new prefix
                    remain = remain[len(prefix2) * prefix_count:]

                if prefix_count == 0:
                    # No valid prefix
                    continue

                suffix_count = count_prefix(remain, suffix)

                if suffix_count == 0:
                    # No valid suffix
                    continue

                # Trim the suffix
                remain = remain[len(suffix) * suffix_count:]
                if len(remain) > 0:
                    # There is something after the suffix
                    continue

                # Make sure we've seen at least the same number of prefixes
                if prefix_count >= suffix_count:
                    return True

        return False


def part2():
    # Update two changed rules
    # 0: 8 11

    # Test short circuits
    # RULE_INDEX[42] = parse_rule('42: "a"')[1]
    # RULE_INDEX[31] = parse_rule('31: "b"')[1]

    # rule8 = parse_rule("8: 42 | 42 8")[1]
    # rule11 = parse_rule("11: 42 31 | 42 11 31")[1]

    # print(RULE_INDEX[8].values())
    # print(RULE_INDEX[11].values())

    rules_parsed = False
    matched = 0

    # Empty rules
    RULE_INDEX.clear()

    with open("input.txt") as f:
        for line in f:
            line = line.strip()
            print(line, end="")
            if line == "":
                rules_parsed = True
                # All rules parsed, replace values
                # RULE_INDEX[8] = rule8
                # RULE_INDEX[11] = rule11
                print("DONE PARSING RULES")
                r42 = RULE_INDEX[42]
                print("prefixes", r42.values())
                r31 = RULE_INDEX[31]
                print("suffixes", r31.values())
            elif not rules_parsed:
                key, rule = parse_rule(line)
                RULE_INDEX[key] = rule
                print(" Parsed")
            elif Rule0().validate(line):
                matched += 1
                print(" OK")
            else:
                print(" Not OK")

    print(f"Total matched is {matched}")


def test2():
    r0 = Rule0()
    # RULE_INDEX[42] = ValueRule(42, ["prefixa", "prefixb"])
    # RULE_INDEX[31] = ValueRule(31, ["sa", "sb"])

    test_input = (
        "42: 9 14 | 10 1",
        "9: 14 27 | 1 26",
        "10: 23 14 | 28 1",
        '1: "a"',
        "11: 42 31",
        "5: 1 14 | 15 1",
        "19: 14 1 | 14 14",
        "12: 24 14 | 19 1",
        "16: 15 1 | 14 14",
        "31: 14 17 | 1 13",
        "6: 14 14 | 1 14",
        "2: 1 24 | 14 4",
        "0: 8 11",
        "13: 14 3 | 1 12",
        "15: 1 | 14",
        "17: 14 2 | 1 7",
        "23: 25 1 | 22 14",
        "28: 16 1",
        "4: 1 1",
        "20: 14 14 | 1 15",
        "3: 5 14 | 16 1",
        "27: 1 6 | 14 18",
        '14: "b"',
        "21: 14 1 | 1 14",
        "25: 1 1 | 1 14",
        "22: 14 14",
        "8: 42",
        "26: 14 22 | 1 20",
        "18: 15 15",
        "7: 14 5 | 1 21",
        "24: 14 1",
    )
    RULE_INDEX.update({
        parse_rule(line)
        for line in test_input
    })

    og_valid = 0
    new_valid = 0

    for v in (
        "abbbbbabbbaaaababbaabbbbabababbbabbbbbbabaaaa",
        "bbabbbbaabaabba",
        "babbbbaabbbbbabbbbbbaabaaabaaa",
        "aaabbbbbbaaaabaababaabababbabaaabbababababaaa",
        "bbbbbbbaaaabbbbaaabbabaaa",
        "bbbababbbbaaaaaaaabbababaaababaabab",
        "ababaaaaaabaaab",
        "ababaaaaabbbaba",
        "baabbaaaabbaaaababbaababb",
        "abbbbabbbbaaaababbbbbbaaaababb",
        "aaaaabbaabaaaaababaa",
        "aaaabbaaaabbaaa",
        "aaaabbaabbaaaaaaabbbabbbaaabbaabaaa",
        "babaaabbbaaabaababbaabababaaab",
        "aabbbbbaabbbaaaaaabbbbbababaaaaabbaaabba",
    ):
        og = RULE_INDEX[0].validate(v)
        n = r0.validate(v)
        print(v, "og", og, "new", n)
        if og:
            og_valid += 1
        if n:
            new_valid += 1

    print("og valid:", og_valid, "new valid:", new_valid)


if __name__ == "__main__":
    # part1()
    part2()
    # test2()