from __future__ import annotations
from dataclasses import dataclass
import bisect
from typing import Optional


import regex as re

from aoc import Aoc2, AocSameParser


class Map:
    def __init__(self, group: list[str]) -> None:
        self.maps: list[tuple[int, int, int]] = []

        for line in group:
            start_dest, start_source, id_range = [
                int(d) for d in re.findall(r"\d+", line)
            ]
            bisect.insort(
                self.maps,
                (start_source, start_dest - start_source, start_source + id_range - 1),
            )
        zero_offsets = []
        start = 0
        for section_start, _, section_end in self.maps:
            if section_start > start:
                zero_offsets.append((start, 0, start - 1))
            start = section_end + 1
        for m in zero_offsets:
            bisect.insort(self.maps, m)

    def find_bucket(self, x: int) -> Optional[int]:
        for idx, (source, _, end) in enumerate(self.maps):
            if source <= x <= end:
                return idx
        return None

    def __getitem__(self, x: int) -> int:
        for source, offset, end in self.maps:
            if source <= x <= end:
                return x + offset
        return x

    def __repr__(self) -> str:
        map_strs = []
        for k, offset, end in self.maps:
            map_strs.append(f"{k} - {end}: {k+offset} - {k+ offset + end}")
        return f"Map({', '.join(map_strs)})"


@dataclass
class Problem:
    seeds: list[int]
    maps: dict[tuple[str, str], Map]

    def convert(self, cat_from: str, cat_to: str, value: int) -> int:
        target = cat_from
        output_value = value
        while target != cat_to:
            new_target = list(
                filter(lambda fromto: fromto[0] == target, self.maps.keys())
            )[0][1]
            output_value = self.maps[(target, new_target)][output_value]
            target = new_target
        return output_value

    @staticmethod
    def parse_header(line: str) -> tuple[str, str]:
        header_match = re.match(r"(?P<source>[^-]*)-to-(?P<dest>[^ ]*) map:", line)
        if header_match is None:
            raise RuntimeError(f'Could not parse "{line}" as group header')

        return header_match["source"], header_match["dest"]

    @staticmethod
    def from_input(inpt: str) -> Problem:
        lines = inpt.splitlines()
        seed_match = re.match(r"seeds:(\s*(?P<seedid>\d+)\s*)*", lines[0])
        if seed_match is None:
            raise RuntimeError(f'Could not parse "{lines[0]}" as seed list')
        seeds = [int(s) for s in seed_match.captures("seedid")]

        groups: list[list[str]] = []
        for line in lines[1:]:
            if not line:
                groups.append([])
            else:
                groups[-1].append(line)
        maps = {Problem.parse_header(group[0]): Map(group[1:]) for group in groups}
        return Problem(seeds, maps)


class Day05(AocSameParser[Problem], Aoc2[Problem, Problem]):
    def __init__(self) -> None:
        Aoc2.__init__(self, 2023, 5, example_code_nr2=0)

    def parseinput(self, inpt: str) -> Problem:
        return Problem.from_input(inpt)

    def part1(self, inpt: Problem) -> int:
        return min(inpt.convert("seed", "location", x) for x in inpt.seeds)

    def part2(self, inpt: Problem) -> int:
        order = [
            ("seed", "soil"),
            ("soil", "fertilizer"),
            ("fertilizer", "water"),
            ("water", "light"),
            ("light", "temperature"),
            ("temperature", "humidity"),
            ("humidity", "location"),
        ]
        seed_iter = iter(inpt.seeds)
        seeds = [
            (start, start + rnge - 1)
            for (start, rnge) in list(zip(seed_iter, seed_iter))
        ]
        for translation in order:
            new_to_convert = []
            current_map = inpt.maps[translation]
            for seed in seeds:
                cur_start, cur_end = seed
                index = current_map.find_bucket(cur_start)
                if index is not None:
                    section_start, section_offset, section_end = current_map.maps[index]
                    while cur_end > section_end:
                        if index == len(current_map.maps) - 1:
                            new_to_convert.append((section_end + 1, cur_end))
                            break
                        new_to_convert.append(
                            (cur_start + section_offset, section_end + section_offset)
                        )
                        index += 1
                        section_start, section_offset, section_end = current_map.maps[
                            index
                        ]
                        cur_start = section_start

                    new_to_convert.append(
                        (cur_start + section_offset, cur_end + section_offset)
                    )
                else:
                    new_to_convert.append(seed)
            seeds = new_to_convert

        return min(x[0] for x in seeds)


def main() -> None:
    day5 = Day05()
    day5.run()


if __name__ == "__main__":
    main()