from enum import Enum
from tools.aoc import AOCDay
from typing import Any


class Material(Enum):
    ORE = 1
    CLAY = 2
    OBSIDIAN = 3
    GEODE = 4


class Robot:
    def __init__(self, produces: Material, costs: dict, min_viable: int = 0):
        self.produces = produces
        self.costs = costs
        self.min_viable = min_viable


class Blueprint:
    def __init__(self, bp_id: int, robots: dict):
        self.bp_id = bp_id
        self.robots = robots


class Inventory:
    def __init__(self):
        self.robots = {
            Material.ORE: 1,
            Material.CLAY: 0,
            Material.OBSIDIAN: 0,
            Material.GEODE: 0,
        }
        self.materials = {
            Material.ORE: 0,
            Material.CLAY: 0,
            Material.OBSIDIAN: 0,
            Material.GEODE: 0,
        }
        self.__build_queue = []

    def gather(self) -> None:
        for mat, amount in self.robots.items():
            self.materials[mat] += amount

    def queue_build_robot(self, robot: Robot) -> None:
        for cost_mat, cost_amount in robot.costs.items():
            self.materials[cost_mat] -= cost_amount
        self.__build_queue.append(robot)

    def build_queued_robots(self) -> None:
        for robot in self.__build_queue:
            self.robots[robot.produces] += 1
        self.__build_queue = []

    def build_robot(self, robot: Robot) -> None:
        self.queue_build_robot(robot)
        self.build_queued_robots()

    def copy(self) -> 'Inventory':
        new_inv = Inventory()
        new_inv.robots = self.robots.copy()
        new_inv.materials = self.materials.copy()
        return new_inv


def pass_minute(bp: Blueprint, inv: Inventory, ignore: list) -> (Inventory, list):
    could_build = []
    for mat in Material:
        for robot_mat, robot_cost in bp.robots[mat].costs.items():
            if inv.materials[robot_mat] < robot_cost:
                break
        else:
            if bp.robots[mat] not in ignore:
                could_build.append(bp.robots[mat])

    inv.gather()

    return inv, could_build


def get_quality(bp: Blueprint, inv: Inventory, ignore: list, time_left: int = 24) -> int:
    for m in range(time_left):
        inv, could_build = pass_minute(bp, inv, ignore)

        if could_build and time_left - m > 1:
            max_quality = 0
            build_a_bot = False
            for robot in could_build:
                if robot.min_viable > time_left - m:
                    continue
                sub_inv = inv.copy()
                sub_inv.build_robot(robot)
                sub_quality = get_quality(bp, sub_inv, [], time_left - m - 1)
                if sub_quality > max_quality:
                    max_quality = sub_quality
                build_a_bot = True

            if build_a_bot:
                sub_quality = get_quality(bp, inv, could_build, time_left - m - 1)
                if sub_quality > max_quality:
                    max_quality = sub_quality

                return max_quality

    return inv.materials[Material.GEODE]


class Day(AOCDay):
    inputs = [
        [
            (33, "input19_test"),
            (1616, "input19"),
        ],
        [
            (56*62, "input19_test"),
            (None, "input19"),
        ]
    ]

    def get_blueprints(self) -> list:
        bp = []
        for line in self.getInput():
            parts = line.split(" ")
            blueprint_id = int(parts[1][:-1])
            ore_ore_cost = int(parts[6])
            clay_ore_cost = int(parts[12])
            obsi_ore_cost = int(parts[18])
            obsi_clay_cost = int(parts[21])
            geode_ore_cost = int(parts[27])
            geode_obsi_cost = int(parts[30])
            robots = {
                Material.ORE: Robot(Material.ORE, {Material.ORE: ore_ore_cost}, 7),
                Material.CLAY: Robot(Material.CLAY, {Material.ORE: clay_ore_cost}, 5),
                Material.OBSIDIAN: Robot(Material.OBSIDIAN, {Material.ORE: obsi_ore_cost, Material.CLAY: obsi_clay_cost}, 3),
                Material.GEODE: Robot(Material.GEODE, {Material.ORE: geode_ore_cost, Material.OBSIDIAN: geode_obsi_cost}, 1)
            }
            bp.append(Blueprint(blueprint_id, robots))

        return bp

    def part1(self) -> Any:
        blueprints = self.get_blueprints()
        score = 0
        for b in blueprints:
            quality = get_quality(b, Inventory(), [])
            print("BP", b.bp_id, "Q", quality)
            score += quality * b.bp_id
        return score

    def part2(self) -> Any:
        blueprints = self.get_blueprints()
        score = 1
        for b in blueprints[:3]:
            quality = get_quality(b, Inventory(), [], 32)
            print("BP", b.bp_id, "Q", quality)
            score *= quality
        return score


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