# -*- coding: utf-8 -*-
import json
from abc import ABC
from functools import cmp_to_key, reduce
from typing import Iterator, Any, List, Tuple, Union

from aoc import BaseAssignment


Packet = Union[List[int], int, List["Packet"]]


class Assignment(BaseAssignment, ABC):
    def compare_pair(self, pair: Tuple[Packet, Packet]):
        packet_a, packet_b = pair

        if isinstance(packet_a, int) and isinstance(packet_b, int):
            diff = packet_b - packet_a
            return diff / abs(diff) if diff != 0 else 0

        if isinstance(packet_a, list) and isinstance(packet_b, int):
            packet_b = [packet_b]

        if isinstance(packet_b, list) and isinstance(packet_a, int):
            packet_a = [packet_a]

        try:
            for sub_pair in zip(packet_a, packet_b, strict=True):
                comparison = self.compare_pair(sub_pair)

                if comparison != 0:
                    return comparison
        except ValueError:
            return self.compare_pair((len(packet_a), len(packet_b)))

        return 0


class AssignmentOne(Assignment):
    example_result = 13

    def pairs(self, input: Iterator[str]):
        while True:
            yield (json.loads(next(input)), json.loads(next(input)))

            try:
                next(input)
            except StopIteration:
                break

    def run(self, input: Iterator) -> Any:
        return sum(
            index
            for index, pair in enumerate(self.pairs(input), start=1)
            if self.compare_pair(pair) > 0
        )


class AssignmentTwo(Assignment):
    example_result = 140

    def packets(self, input: Iterator[str]):
        for item in input:
            if item == "":
                continue

            yield json.loads(item)

    def run(self, input: Iterator) -> Any:
        packets = sorted(
            [[[2]], [[6]], *self.packets(input)],
            key=cmp_to_key(lambda a, b: self.compare_pair((a, b))),
            reverse=True,
        )

        return reduce(
            lambda o, i: o * i,
            [
                index
                for index, packet in enumerate(packets, start=1)
                if packet
                in (
                    [[2]],
                    [[6]],
                )
            ],
        )