import math
from abc import ABC
from dataclasses import dataclass
from typing import List, Tuple, Iterator, TypedDict, Dict, Union

from aoc import BaseAssignment

@dataclass
class Packet:
    raw: str
    version: int = None
    type: int = None
    length_type = str
    padding: int = 0
    payload_length = None
    payload_start = None
    payload: Union[int, List['Packet']] = None

    LITERAL = 4

    FIFTEEN = '0'
    ELEVEN = '1'

    PAYLOAD_LENGTH_TYPE = {
        FIFTEEN: 15,
        ELEVEN: 11
    }

    @classmethod
    def from_hex(cls, value: str) -> 'Packet':
        print()
        print(value)
        for byte in value:
            print(byte, bin(int(byte, 16))[2:].rjust(4, '0'))

        bits = ''.join([
            bin(int(byte, 16))[2:]
                .rjust(4, '0')
            for byte
            in value
        ])

        return cls(
            raw=bits
        )


    def __post_init__(self):
        self.version = int(self.raw[:3], 2)
        self.type = int(self.raw[3:6], 2)

        match self.type:
            case self.LITERAL:
                self.payload = self.parse_literal()
            case _:
                self.length_type = self.raw[6]
                self.payload_start = 8 + self.PAYLOAD_LENGTH_TYPE[self.length_type]
                self.payload_length = int(self.raw[8:self.payload_start], 2)

                print(self.raw[:3], self.raw[3:6], self.raw[7], self.raw[7:self.payload_start], self.raw[self.payload_start:])

                self.payload = self.parse_operator()


    def __len__(self):
        return len(self.raw) - self.padding

    def pad(self):
        added_zeroes = 4 - len(self.raw) % 4
        self.raw += added_zeroes * '0'
        self.padding = added_zeroes

    def parse_literal(self) -> int:
        start_index = 6

        end = False
        bits = []

        while not end:
            bits.append(self.raw[start_index + 1:start_index + 5])
            if self.raw[start_index] == '0':
                end = True

            start_index += 5

        self.raw = self.raw[:start_index]
        self.pad()

        return int(''.join(bits), 2)

    def parse_operator(self) -> List['Packet']:
        packets = []
        payload_end = self.payload_start + self.payload_length
        payload = self.raw[self.payload_start:payload_end]

        while len(payload) > 0:
            packets.append(Packet(raw=payload))
            payload = payload[len(packets[-1]):]

        self.raw = self.raw[:payload_end]
        self.pad()

        return packets


class Assignment(BaseAssignment, ABC):
    def parse_item(self, item: str) -> Packet:
        return Packet.from_hex(item)


    def read_input(self, example=False) -> Packet:
        return next(super().read_input(example))

class AssignmentOne(Assignment):
    example_result = 31

    def version_sum(self, packet: Packet):
        if packet.type == Packet.LITERAL:
            return packet.version

        return packet.version + sum(
            self.version_sum(p)
            for p in packet.payload
        )

    def run(self, input: Packet):
        return self.version_sum(input)