from tools.aoc import AOCDay
from typing import Any


ROCK_SHAPES = [
    ["####"],
    [' # ', '###', ' # '],
    ['  #', '  #', '###'],
    ['#', '#', '#', '#'],
    ['##', '##']
]


class Chamber(list):
    def hash(self):
        max_len = min(30, len(self) - 1)
        return "".join(self[i] for i in range(max_len))


def collide_hor(chamber: list, rock_shape:list, rock_x: int, rock_y: int) -> bool:
    if rock_x < 0 or rock_x > 7 - len(rock_shape[0]):
        return True

    for i, line in enumerate(rock_shape):
        for x, c in enumerate(rock_shape[i]):
            if c == ' ':
                continue
            y = rock_y - ((len(rock_shape) - 1) - i)
            if y < 0:
                continue
            if chamber[y][rock_x + x] == '#':
                return True

    return False


def settles(chamber: list, rock_shape: list, rock_x: int, rock_y: int) -> bool:
    for i, c in enumerate(rock_shape[-1]):
        if c == ' ':
            continue
        if rock_y >= -1 and chamber[rock_y + 1][i + rock_x] == '#':
            return True

    if rock_shape[-1] == ' # ':
        for i, c in enumerate(rock_shape[-2]):
            if c == ' ':
                continue
            if rock_y > -1 and chamber[rock_y][i + rock_x] == '#':
                return True

    return False


def drop_rock(chamber: list, rock_shape: list, jet: list, jet_index: int) -> int:
    jet_len = len(jet)
    rock_index_x = 2
    for i in range(4):
        rock_index_x += jet[(jet_index + i) % jet_len]
        if not (0 <= rock_index_x <= 7 - len(rock_shape[0])):
            rock_index_x = 0 if rock_index_x < 0 else 7 - len(rock_shape[0])
    jet_index = (jet_index + 4) % jet_len
    rock_index_y = -1
    while not settles(chamber, rock_shape, rock_index_x, rock_index_y):
        rock_index_y += 1
        if not collide_hor(chamber, rock_shape, rock_index_x + jet[jet_index], rock_index_y):
            rock_index_x += jet[jet_index]
        jet_index = (jet_index + 1) % jet_len

    for i in range(len(rock_shape) - 1, -1, -1):
        y = rock_index_y + i - len(rock_shape) + 1
        if y < 0:
            chamber.insert(0, '       ')
            y = 0
        left = chamber[y][:rock_index_x]
        middle = "".join(
            "#" if rock_shape[i][x] == '#' or chamber[y][rock_index_x + x] == '#' else ' '
            for x in range(len(rock_shape[i]))
        )
        right = chamber[y][rock_index_x + len(rock_shape[i]):]
        chamber[y] = left + middle + right

    return jet_index


class Day(AOCDay):
    inputs = [
        [
            (3068, "input17_test"),
            (3059, "input17"),
        ],
        [
            (1514285714288, "input17_test"),
            (1500874635587, "input17"),
        ]
    ]

    def get_jet_pattern(self) -> list:
        return [-1 if x == '<' else 1 for x in self.getInput()]

    def part1(self) -> Any:
        chamber = ['#######']
        jet = self.get_jet_pattern()
        jet_index = 0
        for rock_fall in range(2022):
            jet_index = drop_rock(chamber, ROCK_SHAPES[rock_fall % 5], jet, jet_index)

        return len(chamber) - 1

    def part2(self) -> Any:
        chamber = Chamber(['#######'])
        jet = self.get_jet_pattern()
        jet_index = drop_rock(chamber, ROCK_SHAPES[0], jet, 0)
        rock_index = 1
        cycle = 1
        cycle_cache = {}
        while (jet_index, rock_index, chamber.hash()) not in cycle_cache:
            cycle_cache[(jet_index, rock_index, chamber.hash())] = (cycle, len(chamber) - 1)
            jet_index = drop_rock(chamber, ROCK_SHAPES[rock_index], jet, jet_index)
            rock_index = (rock_index + 1) % 5
            cycle += 1

        cycle_start, cycle_start_len = cycle_cache[(jet_index, rock_index, chamber.hash())]
        cycle_height = len(chamber) - 1 - cycle_start_len
        cycle_len = cycle - cycle_start

        chamber = Chamber(['#######'])
        jet_index = 0
        rock_index = 0
        bottom_part = 1_000_000_000_000 % cycle_len
        for _ in range(cycle_len + bottom_part):
            jet_index = drop_rock(chamber, ROCK_SHAPES[rock_index % 5], jet, jet_index)
            rock_index += 1
        cur_len = len(chamber) - 1 - cycle_height

        return (1_000_000_000_000 // cycle_len * cycle_height) + cur_len


if __name__ == '__main__':
    day = Day(2022, 17)
    day.run(verbose=True)