#!/usr/bin/env python3
from pathlib import Path
from math import prod


def get_next_packet(msg, idx, packets, res=0):
    """The function is called recursively and will read a single packet starting at index 'idx' in
    'msg'. The packet will be appended to the 'packets' list. For each packet the type_id and a
    payload is stored. The payload is a literal for packets of type 4 or will be filled with
    subsequent packages for other types by recursion of this function. The version will be added to
    the result.
    At the end the updated index and result will be returned.
    """
    if (idx+11) >= len(msg):
        idx = -1
        return idx, res
    version = int(msg[idx:idx + 3], 2)
    idx += 3
    res += version
    type_id = int(msg[idx:idx + 3], 2)
    idx += 3
    payload = None
    match type_id:
        case 4:
            end = False
            literal = 0
            while not end:
                if not int(msg[idx:idx + 1], 2):
                    end = True
                idx += 1
                literal = (literal << 4) | int(msg[idx:idx + 4], 2)
                idx += 4
            payload = literal
        case _:
            length_id = int(msg[idx:idx + 1], 2)
            idx += 1
            length = 0
            payload = []
            if length_id:
                length = int(msg[idx:idx + 11], 2)
                idx += 11
                for _ in range(length):
                    idx, res = get_next_packet(msg, idx, payload, res)
            else:
                length = int(msg[idx:idx + 15], 2)
                idx += 15
                next_idx = idx + length
                while idx < next_idx:
                    idx, res = get_next_packet(msg, idx, payload, res)
    packets.append({'type_id': type_id, 'payload': payload})
    return idx, res


def calc_result(packets):
    """The function recursively steps through the packets and performs the defined calculations."""
    res = 0
    for p in packets:
        match p['type_id']:
            case 4:
                res = p['payload']
            case _:
                sub_packets = []
                for i in range(len(p['payload'])):
                    sub_packets.append(calc_result([p['payload'][i]]))
                match p['type_id']:
                    case 0:
                        res = sum(sub_packets)
                    case 1:
                        res = prod(sub_packets)
                    case 2:
                        res = min(sub_packets)
                    case 3:
                        res = max(sub_packets)
                    case 5:
                        res = 1 if sub_packets[0] > sub_packets[1] else 0
                    case 6:
                        res = 1 if sub_packets[0] < sub_packets[1] else 0
                    case 7:
                        res = 1 if sub_packets[0] == sub_packets[1] else 0
    return res


def solve(input):
    """Part 1 and part 2 are solved in one go so we don't have to parse the bitstream twice. We
    start by stripping newline and zeros at the end of the input line, convert it to an integer
    with base 16 (hexadecimal) and back to a binary string. The bin() function adds '0b' at the
    beginning so we only store te actual binary data. Afterwards we need to add zeros in the
    beginning until the binary string has a length of four times the hexadecimal representation
    (Each hex digit is four binary digits).
    """
    result_p1 = 0
    line = input[0].rstrip().strip('0')
    msg = bin(int(line, 16))[2:]
    while len(msg) < len(line) * 4:
        msg = '0' + msg
    packets = []
    next = 0
    while -1 != next:
        next, result_p1 = get_next_packet(msg, next, packets, result_p1)
    result_p2 = calc_result(packets)
    print("Part 1 result:", result_p1)
    print("Part 2 result:", result_p2)


input = list()
p = Path(__file__).with_name('input.txt')
with open(p) as f:
    input = f.readlines()
    solve(input)