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):
        self.produces = produces
        self.costs = costs


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) -> 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 could_build


def get_quality(bp: Blueprint, inv: Inventory, ignore: list, time_left: int = 24, min_geode: int = 0) -> (int, int):
    #print("get_quality", time_left, inv.materials, inv.robots)
    for m in range(time_left):
        could_build = pass_minute(bp, inv, ignore)
        if bp.robots[Material.GEODE] in could_build and time_left - m > min_geode:
            min_geode = time_left - m
            print("Can build Geode Robot with", min_geode, "min left.")
            print(inv.robots, inv.materials)

        if could_build and time_left - m > 1 and (inv.robots[Material.GEODE] > 0 or time_left > min_geode):
            max_quality = 0
            for robot in could_build:
                sub_inv = inv.copy()
                sub_inv.build_robot(robot)
                sub_quality, sub_min_geode = get_quality(bp, sub_inv, [], time_left - m - 1, min_geode)
                if sub_quality > max_quality:
                    max_quality = sub_quality
                if sub_min_geode > min_geode:
                    min_geode = sub_min_geode

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

            return max_quality, min_geode

    return inv.materials[Material.GEODE], min_geode


class Day(AOCDay):
    inputs = [
        [
            (33, "input19_test"),
            (None, "input19"),
        ],
        [
            (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}),
                Material.CLAY: Robot(Material.CLAY, {Material.ORE: clay_ore_cost}),
                Material.OBSIDIAN: Robot(Material.OBSIDIAN, {Material.ORE: obsi_ore_cost, Material.CLAY: obsi_clay_cost}),
                Material.GEODE: Robot(Material.GEODE, {Material.ORE: geode_ore_cost, Material.OBSIDIAN: geode_obsi_cost})
            }
            bp.append(Blueprint(blueprint_id, robots))

        return bp

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

    def part2(self) -> Any:
        return ""


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