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


def get_next_packet(msg, start, packets, res=0):
    idx = start
    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:
            # literal
            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
            # print('l', version, type_id, 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
                # print('o', version, type_id, length_id, length)
                for _ in range(length):
                    idx, res = get_next_packet(msg, idx, payload, res)
            else:
                length = int(msg[idx:idx + 15], 2)
                idx += 15
                # print('o', version, type_id, length_id, length)
                next_idx = idx + length
                while idx < next_idx:
                    idx, res = get_next_packet(msg, idx, payload, res)
    packets.append({'version': version, 'type_id': type_id, 'payload': payload})
    return idx, res


def calc_result(packets):
    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 part_1(input):
    result = 0
    line = input[0].rstrip().strip('0')
    msg = bin(int(line, 16))
    msg = msg[2:]
    while len(msg) % 4:
        msg = '0' + msg
    packets = []
    next = 0
    while -1 != next:
        next, result = get_next_packet(msg, next, packets, result)
    print("Part 1 result:", result)


def part_2(input):
    result = 0
    line = input[0].rstrip().strip('0')
    msg = bin(int(line, 16))
    msg = msg[2:]
    while len(msg) % 4:
        msg = '0' + msg
    packets = []
    next = 0
    while -1 != next:
        next, _ = get_next_packet(msg, next, packets)
    result = calc_result(packets)
    print("Part 2 result:", result)


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