import itertools
from dataclasses import dataclass
from typing import List
import pathlib


def read_data():
    with open(f"data/{pathlib.Path(__file__).stem}.txt") as raw_data:
        return [line.strip() for line in raw_data.readlines()]


Data = List['Data'] | int
Stream = List[Data]


@dataclass
class Packet:
    stream: Stream

    def __lt__(self, other: 'Packet'):
        return test_data_order(self.stream, other.stream) == -1


def find_close_bracket(line: str, start: int) -> int:
    """ find the index of the closing bracket of the first, outermost bracket pair"""
    bracket_count = 0
    for idx, ch in itertools.islice(enumerate(line), start, None):
        match ch:
            case '[':
                bracket_count += 1
            case ']':
                bracket_count -= 1
                if bracket_count == 0:
                    return idx


def parse_stream(line: str) -> Stream:
    stream: Stream = []
    idx = 1

    # ignore surrounding brackets
    while idx < len(line) - 1:
        if line[idx] == '[':
            close_idx = find_close_bracket(line, idx)
            stream.append(parse_stream(line[idx:close_idx + 1]))
            idx = close_idx + 2

        else:
            try:
                next_idx = line[idx:].index(",")
                stream.append(int(line[idx:idx + next_idx]))
                idx += next_idx + 1
            except Exception as e:
                # we're on the last number
                stream.append(int(line[idx:-1]))
                break

    return stream


def test_data_order(left: Data, right: Data) -> int:
    match left, right:
        case [], []:
            return 0
        case [_, *_], []:
            return 1
        case [], [_, *_]:
            return -1
        case [x, *xs], [y, *ys]:
            if (ret := test_data_order(x, y)) != 0:
                return ret
            return test_data_order(xs, ys)
        case int(x), int(y):
            if x < y:
                return -1
            elif x > y:
                return 1
            else:
                return 0
        case int(x), [*y]:
            return test_data_order([x], y)
        case [*x], int(y):
            return test_data_order(x, [y])


def parse_packets(data: List[str]) -> List[Packet]:
    packets: List[Packet] = []
    data = iter(data)

    while True:

        left = Packet(parse_stream(next(data)))
        right = Packet(parse_stream(next(data)))
        packets.extend((left, right))

        try:
            next(data)
        except:
            break

    return packets


def part1(data: List[str]) -> int:
    packets = parse_packets(data)
    order_sum = 0
    for idx, (left, right) in enumerate(zip(packets[::2], packets[1::2])):
        if left < right:
            order_sum += idx + 1

    return order_sum


def part2(data: List[str]) -> int:
    divider1 = Packet(parse_stream("[[2]]"))
    divider2 = Packet(parse_stream("[[6]]"))

    packets = parse_packets(data)
    packets.extend([divider1, divider2])

    packets.sort()

    return (packets.index(divider1) + 1) * (packets.index(divider2) + 1)


def main():
    data = read_data()
    print(f"Part 1: {part1(data)}")
    print(f"Part 2: {part2(data)}")


if __name__ == "__main__":
    main()