from __future__ import annotations
from dataclasses import dataclass
from enum import IntEnum
from itertools import islice
from math import prod
from multiprocessing import Pool
from queue import PriorityQueue

from typing import Iterable, Iterator, Self
from advent.parser.parser import P

day_num = 19


def part1(lines: Iterator[str]) -> int:
    return sum(bp.number * geodes for bp, geodes in Processor.pool_it(24, lines))


def part2(lines: Iterator[str]) -> int:
    return prod(num for _, num in Processor.pool_it(32, islice(lines, 3)))


@dataclass(slots=True, frozen=True)
class Processor:
    rounds: int

    @classmethod
    def pool_it(cls, rounds: int, lines: Iterator[str]) -> Iterable[tuple[Blueprint, int]]:
        with Pool() as p:
            return p.map(Processor(rounds), lines)

    def __call__(self, line: str) -> tuple[Blueprint, int]:
        blueprint = Blueprint.parse(line)
        return blueprint, blueprint.run(self.rounds)


number_parser = P.second(P.string("Blueprint "), P.unsigned())
ore_parser = P.second(P.string(": Each ore robot costs "), P.unsigned())
clay_parser = P.second(P.string(" ore. Each clay robot costs "), P.unsigned())
tuple_parser = P.seq(P.unsigned(), P.second(P.string(" ore and "), P.unsigned()))
obsidian_parser = P.second(P.string(" ore. Each obsidian robot costs "), tuple_parser)
geode_parser = P.second(P.string(" clay. Each geode robot costs "), tuple_parser)
blueprint_parser = P.map5(number_parser, ore_parser, clay_parser, obsidian_parser, geode_parser,
                          lambda number, ore, clay, obsidian, geode:
                          Blueprint.create(number, ore, clay, obsidian, geode))


class Element(IntEnum):
    Geode = 0
    Obsidian = 1
    Clay = 2
    Ore = 3


Elements = tuple[int, int, int, int]


def gt(first: Elements, second: Elements) -> bool:
    return all(f >= s for f, s in zip(first, second))


def add(first: Elements, second: Elements) -> Elements:
    return tuple(v1 + v2 for v1, v2 in zip(first, second))


def sub(first: Elements, second: Elements) -> Elements:
    return Elements(tuple(v1 - v2 for v1, v2 in zip(first, second)))


def inc_element(elements: Elements, pos: Element) -> Elements:
    return tuple(v + 1 if num == pos else v for num, v in enumerate(elements))


@dataclass(slots=True)
class Path:
    time: int
    material: Elements
    robots: Elements
    blueprint: Blueprint
    path: list[Element | None]

    @classmethod
    def start(cls, blueprint: Blueprint) -> Path:
        return Path(0, (0, 0, 0, 0), (0, 0, 0, 1), blueprint, [])

    def _check(self, element: Element) -> Path | None:
        if element != Element.Geode:
            max_needed = max(req[element] for req in self.blueprint.requirements)
            if self.robots[element] >= max_needed:
                return None

        if gt(self.material, self.blueprint.requirements[element]):
            return Path(self.time + 1,
                        add(sub(self.material, self.blueprint.requirements[element]), self.robots),
                        inc_element(self.robots, element), self.blueprint, self.path + [element])
        else:
            return None

    def find_next(self) -> Iterator[Path]:
        if (path := self._check(Element.Geode)) is not None:
            yield path

        if self.blueprint.max_requirement(Element.Obsidian) >= self.material[Element.Obsidian]:
            if (path := self._check(Element.Obsidian)) is not None:
                yield path

        if self.blueprint.max_requirement(Element.Clay) >= self.material[Element.Clay]:
            if (path := self._check(Element.Clay)) is not None:
                yield path

        if self.blueprint.max_requirement(Element.Ore) >= self.material[Element.Ore]:
            if (path := self._check(Element.Ore)) is not None:
                yield path

        yield Path(self.time + 1, add(self.material, self.robots), self.robots, self.blueprint,
                   self.path + [None])

    def __lt__(self, other: Path) -> bool:
        if self.time != other.time:
            return self.time < other.time
        return self.material > other.material


@dataclass(slots=True)
class Blueprint:
    number: int
    requirements: tuple[Elements, Elements, Elements, Elements]

    @classmethod
    def create(cls, number: int, ore: int, clay: int,
               obsidian: tuple[int, int], geode: tuple[int, int]) -> Self:
        requirements = (
            (0, geode[1], 0, geode[0]),
            (0, 0, obsidian[1], obsidian[0]),
            (0, 0, 0, clay),
            (0, 0, 0, ore),
        )
        return Blueprint(number, requirements)

    @classmethod
    def parse(cls, line: str) -> Self:
        return blueprint_parser.parse(line).get()

    def max_requirement(self, element: Element) -> int:
        return round(max(requirement[element] for requirement in self.requirements) * 1.2)

    def run(self, rounds: int) -> int:
        queue: PriorityQueue[Path] = PriorityQueue()
        queue.put(Path.start(self))
        seen: dict[tuple[Elements, int], Elements] = {}
        while not queue.empty():
            current = queue.get()
            last_seen = seen.get((current.robots, current.time))
            if last_seen is not None and gt(last_seen, current.material):
                continue
            seen[(current.robots, current.time)] = current.material

            if current.time == rounds:
                return current.material[Element.Geode]
            for next in current.find_next():
                queue.put(next)

        raise Exception("No optimum found")