from math import prod

from tools.aoc import AOCDay
from typing import Any, List


def get_packet(bit_rep: str) -> (str, int):  # packet, packet_length
    msg_type = int(bit_rep[3:6], 2)
    index = 6

    if msg_type == 4:
        while bit_rep[index] == "1":
            index += 5
        index += 5
    else:
        op_type = bit_rep[index]
        index += 1
        if op_type == "0":
            sub_pkg_len = int(bit_rep[index:index+15], 2)
            index += 15 + sub_pkg_len
        else:
            sub_pkg_count = int(bit_rep[index:index+11], 2)
            index += 11
            for x in range(sub_pkg_count):
                _, sub_pkg_len = get_packet(bit_rep[index:])
                index += sub_pkg_len

    return bit_rep[:index], index


def get_subpackages(packet: str) -> List[str]:
    ptype = int(packet[3:6], 2)
    if ptype == 4:
        return [packet]
    else:
        plist = []
        op_type = packet[6]
        if op_type == "0":
            sub_pkg_len = int(packet[7:22], 2)
            sub_packages = packet[22:22+sub_pkg_len]
            while "1" in sub_packages:
                sub_packet, sub_pkg_len = get_packet(sub_packages)
                sub_packages = sub_packages[sub_pkg_len:]
                plist.append(sub_packet)
        else:
            sub_pkg_count = int(packet[7:18], 2)
            index = 18
            for x in range(sub_pkg_count):
                sub_packet, sub_pkg_len = get_packet(packet[index:])
                index += sub_pkg_len
                plist.append(sub_packet)

        return plist


def get_versions(packet: str) -> int:
    version = int(packet[0:3], 2)
    ptype = int(packet[3:6], 2)
    if ptype == 4:
        return version
    else:
        for p in get_subpackages(packet):
            version += get_versions(p)

    return version


def get_value(packet: str) -> int:
    while "1" in packet:
        this_packet, pkg_len = get_packet(packet)
        pkg_type = int(this_packet[3:6], 2)
        packet = packet[pkg_len:]
        if pkg_type == 0:
            return sum(get_value(p) for p in get_subpackages(this_packet))
        elif pkg_type == 1:
            return prod(get_value(p) for p in get_subpackages(this_packet))
        elif pkg_type == 2:
            return min(get_value(p) for p in get_subpackages(this_packet))
        elif pkg_type == 3:
            return max(get_value(p) for p in get_subpackages(this_packet))
        elif pkg_type == 4:
            index = 6
            value = 0
            while this_packet[index] == "1":
                value = value * 16 + int(this_packet[index + 1:index + 5], 2)
                index += 5
            value = value * 16 + int(this_packet[index + 1:index + 5], 2)
            return value
        elif pkg_type == 5:
            sub_pkg = get_subpackages(this_packet)
            return get_value(sub_pkg[0]) > get_value(sub_pkg[1])
        elif pkg_type == 6:
            sub_pkg = get_subpackages(this_packet)
            return get_value(sub_pkg[0]) < get_value(sub_pkg[1])
        elif pkg_type == 7:
            sub_pkg = get_subpackages(this_packet)
            return get_value(sub_pkg[0]) == get_value(sub_pkg[1])


class Day(AOCDay):
    test_solutions_p1 = [16, 12, 23, 31, 1007]
    test_solutions_p2 = [3, 54, 7, 9, 1, 0, 0, 1, 834151779165]

    def getBits(self) -> str:
        bits = bin(int(self.getInput(), 16))[2:]
        return "0" * (len(self.getInput()) * 4 - len(bits)) + bits

    def part1(self) -> Any:
        return get_versions(self.getBits())

    def part2(self) -> Any:
        return get_value(self.getBits())