from enum import Enum

from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


DIRECTIONS = [
    Coordinate(0, -1),
    Coordinate(1, 0),
    Coordinate(0, 1),
    Coordinate(-1, 0)
]


class NodeState(Enum):
    CLEAN = 0
    WEAKENED = 1
    INFECTED = 2
    FLAGGED = 3


class Day(AOCDay):
    inputs = [
        [
            (5587, "input22_test"),
            (5460, "input22")
        ],
        [
            (2511944, "input22_test"),
            (2511702, "input22")
        ]
    ]

    def get_init_state(self) -> (Grid, int, Coordinate):
        grid = Grid.from_str(
            "/".join(self.getInput()),
            true_char='#',
            true_value=NodeState.INFECTED,
            default=NodeState.CLEAN
        )

        carrier_direction_index = 0
        cattier_position = Coordinate((len(self.getInput()[0]) - 1) // 2, (len(self.getInput()) - 1) // 2)

        return grid, carrier_direction_index, cattier_position

    def part1(self) -> Any:
        grid, carrier_direction_index, carrier_position = self.get_init_state()
        infect_count = 0
        for _ in range(10_000):
            if grid.get(carrier_position) == NodeState.INFECTED:
                carrier_direction_index = (carrier_direction_index + 1) % 4
                grid.set(carrier_position, NodeState.CLEAN)
            else:
                infect_count += 1
                carrier_direction_index = (carrier_direction_index - 1) % 4
                grid.set(carrier_position, NodeState.INFECTED)

            carrier_position += DIRECTIONS[carrier_direction_index]

        return infect_count

    def part2(self) -> Any:
        grid, carrier_direction_index, carrier_position = self.get_init_state()
        infect_count = 0
        for _ in range(10_000_000):
            match grid.get(carrier_position):
                case NodeState.CLEAN:
                    carrier_direction_index = (carrier_direction_index - 1) % 4
                    grid.set(carrier_position, NodeState.WEAKENED)
                case NodeState.WEAKENED:
                    grid.set(carrier_position, NodeState.INFECTED)
                    infect_count += 1
                case NodeState.INFECTED:
                    carrier_direction_index = (carrier_direction_index + 1) % 4
                    grid.set(carrier_position, NodeState.FLAGGED)
                case NodeState.FLAGGED:
                    carrier_direction_index = (carrier_direction_index + 2) % 4
                    grid.set(carrier_position, NodeState.CLEAN)

            carrier_position += DIRECTIONS[carrier_direction_index]

        return infect_count


if __name__ == '__main__':
    day = Day(2017, 22)
    day.run(verbose=True)