from itertools import islice, pairwise
from typing import NamedTuple

from more_itertools import chunked, sliced

test_input = """
2333133121414131402
"""

test_solution_p1 = 1928
test_solution_p2 = 2858


def solve_p1(puzzle_input: str) -> int:
    blocks = _parse_blocks(puzzle_input)
    compacted_blocks = _compact_blocks(blocks)
    return _hash_blocks(compacted_blocks)


def solve_p2(puzzle_input: str) -> int:
    blocks = _parse_blocks(puzzle_input)
    compacted_blocks = _compact_blocks_no_fragmentation(blocks)
    return _hash_blocks(compacted_blocks)


class Block(NamedTuple):
    id: int
    position: int
    size: int


def _parse_blocks(puzzle_input: str) -> list[Block]:
    blocks: list[Block] = []
    cursor_position = 0
    current_id = 0

    snapshot = map(int, puzzle_input.strip())
    for pair in chunked(snapshot, 2):
        if data_size := pair[0]:
            blocks.append(Block(current_id, cursor_position, data_size))
            cursor_position += data_size
            current_id += 1
        if len(pair) == 2:
            cursor_position += pair[1]

    return blocks


def _compact_blocks(blocks: list[Block]) -> list[Block]:
    assert blocks
    blocks = blocks.copy()

    compacted_blocks: list[Block] = []
    current_position = 0
    moving_block = blocks.pop()

    while blocks:
        next_block = blocks.pop(0)
        free_space = next_block.position - current_position
        while free_space:
            assert moving_block.size > 0
            allocation = min(free_space, moving_block.size)
            compacted_blocks.append(
                Block(
                    moving_block.id,
                    current_position,
                    allocation,
                )
            )
            current_position += allocation
            free_space -= allocation
            if moving_block.size > allocation:
                moving_block = Block(
                    moving_block.id,
                    moving_block.position + allocation,
                    moving_block.size - allocation,
                )
            elif blocks:
                moving_block = blocks.pop()
            else:
                break

        compacted_blocks.append(next_block)
        current_position += next_block.size

    compacted_blocks.append(
        Block(
            moving_block.id,
            current_position,
            moving_block.size,
        )
    )

    return compacted_blocks


def _compact_blocks_no_fragmentation(blocks: list[Block]) -> list[Block]:
    assert blocks
    blocks = blocks.copy()

    compacted_blocks: list[Block] = []

    while len(blocks) > 1:
        moving_block = blocks[-1]
        for (_, a), (i, b) in pairwise(enumerate(blocks)):
            free_space = b.position - a.position - a.size
            if free_space >= moving_block.size:
                modified_block = Block(
                    moving_block.id,
                    a.position + a.size,
                    moving_block.size,
                )
                blocks.pop()
                blocks.insert(i, modified_block)
                break
        else:
            blocks.pop()
            compacted_blocks.append(moving_block)

    compacted_blocks.append(blocks[0])

    return list(sorted(compacted_blocks, key=lambda block: block.position))


def _print_blocks(blocks: list[Block]) -> str:
    snapshot = []
    for a, b in pairwise(blocks):
        snapshot.append(str(a.id) * a.size)
        snapshot.append("." * (b.position - a.position - a.size))
    snapshot.append(str(blocks[-1].id) * blocks[-1].size)


def _hash_blocks(blocks: list[Block]) -> int:
    return sum(
        block.id * pos
        for block in blocks
        for pos in range(block.position, block.position + block.size)
    )