from collections.abc import Generator
from collections.abc import Iterable
from dataclasses import dataclass
from itertools import batched
from itertools import chain
from pathlib import Path


@dataclass
class RangeMapContent:
    dest_start: int
    source_start: int
    length: int

    def __contains__(self, key: int) -> int:
        return self.source_start <= key < self.source_start + self.length

    def __getitem__(self, key: int) -> int:
        # print(f"looking for {key} in {self.dest_start} {self.source_start}, {self.length}")
        if key not in self:
            return key

        delta = self.source_start - self.dest_start

        return key - delta


class RangeMap:
    def __init__(self, *contents: RangeMapContent) -> None:
        self.contents = contents

    def __getitem__(self, key: int) -> int:
        for content in self.contents:
            if key in content:
                return content[key]

        return key


def part1(input: Path, part2=False):
    seeds: Iterable[int] = []
    maps: dict[str, RangeMap] = {}
    map_name: str = ""
    map_contents: list[RangeMapContent] = []
    with input.open() as f:
        for line in f:
            line = line.strip()
            if not seeds:
                seeds = [int(s.strip()) for s in line.partition(":")[2].split()]
                if part2:
                    windows = [
                        [start, start + length] for start, length in batched(seeds, n=2)
                    ]
                    windows.sort(key=lambda x: x[0])
                    merged_windows: list[list[int]] = []
                    for window in windows:
                        if not merged_windows:
                            merged_windows.append(window)
                            continue
                        if window[0] > merged_windows[-1][1]:
                            merged_windows.append(window)
                            continue
                        if window[1] > merged_windows[-1][1]:
                            merged_windows[-1][1] = window[1]
                            continue

                    seeds = chain.from_iterable(
                        range(window[0], window[1]) for window in merged_windows
                    )

                    # seeds = set(chain(*[
                    #     range(start, start+length)
                    #     for start, length in batched(seeds, n=2)
                    # ]))

                # print(f"{len(list(seeds))} seeds")
                continue

            if line == "" and map_name != "":
                # print(f"Creating map {map_name} with ranges {map_contents}")
                maps[map_name] = RangeMap(*map_contents)
                map_name = ""
                map_contents.clear()
                continue

            if map_name == "":
                map_name = line.partition(" ")[0]
                continue

            items = [int(i) for i in line.split()]
            dest_start = items[0]
            source_start = items[1]
            range_length = items[2]

            map_contents.append(RangeMapContent(dest_start, source_start, range_length))

        maps[map_name] = RangeMap(*map_contents)

    # for name, m in maps.items():
    #     print(f"map {name}")
    #     for c in m.contents:
    #         print(f"  {c.dest_start} {c.source_start} {c.length}")

    min_location: int | None = None

    for seed in seeds:
        soil = maps["seed-to-soil"][seed]
        fertilizer = maps["soil-to-fertilizer"][soil]
        water = maps["fertilizer-to-water"][fertilizer]
        light = maps["water-to-light"][water]
        temperature = maps["light-to-temperature"][light]
        humidity = maps["temperature-to-humidity"][temperature]
        location = maps["humidity-to-location"][humidity]

        if min_location is None or location < min_location:
            min_location = location

        print(
            f"Seed {seed} to soil {soil} to fertilizer {fertilizer} to water {water} to light {light} to temperature {temperature} to humidity {humidity} to location {location}"
        )

    return min_location


if __name__ == "__main__":
    input = Path("./sample.txt")
    # input = Path("./input.txt")
    answer1 = part1(input)
    print("part1", answer1)

    answer2 = part1(input, True)
    print("part2", answer2)