import hashlib

from tools.aoc import AOCDay
from typing import Any


class Day(AOCDay):
    inputs = [
        [
            (22728, "input14_test"),
            (23890, "input14"),
        ],
        [
            (22696, "input14"),
        ],
    ]

    def get_hash(self, to_hash: str, p2: bool = False) -> str:
        if to_hash in self.DP:
            return self.DP[to_hash]

        the_hash = hashlib.md5(to_hash.encode()).hexdigest()
        if p2:
            for _ in range(2016):
                the_hash = hashlib.md5(the_hash.encode()).hexdigest()
        self.DP[to_hash] = the_hash
        return the_hash

    def get_key_index(self, p2: bool = False) -> int:
        salt = self.getInput()
        index = -1
        key_count = 0
        while key_count < 64:
            index += 1
            the_hash = self.get_hash(salt + str(index), p2)
            for i in range(len(the_hash) - 2):
                found_triplet = False
                if the_hash[i] == the_hash[i + 1] == the_hash[i + 2]:
                    found_triplet = True
                    for t in range(index + 1, index + 1001):
                        if the_hash[i] * 5 in self.get_hash(salt + str(t), p2):
                            key_count += 1
                            break

                if found_triplet:
                    break

        return index

    def part1(self) -> Any:
        return self.get_key_index()

    def part2(self) -> Any:
        return self.get_key_index(True)


if __name__ == "__main__":
    day = Day(2016, 14)
    day.run(verbose=True)