from __future__ import annotations
from tools.aoc import AOCDay
from typing import Any


class Block:
    file_id: int
    file_length: int
    left: Block | None
    right: Block | None

    def __init__(self, file_id: int, file_length: int, left: Block | None = None):
        self.file_id = file_id
        self.file_length = file_length
        self.left = left
        self.right = None


def get_checksum(head: Block) -> int:
    checksum = 0
    index = head
    count = 0
    while index is not None:
        print((str(index.file_id) + "|") * index.file_length, end="")
        if index.file_id > 0:
            for _ in range(index.file_length):
                checksum += index.file_id * count
                count += 1
        else:
            count += index.file_length
        index = index.right

    print()
    return checksum


class Day(AOCDay):
    inputs = [
        [
            (1928, "input9_test"),
            (6346871685398, "input9"),
        ],
        [
            (2858, "input9_test"),
            (None, "input9"),
        ],
    ]

    def parse_input(self) -> tuple[Block, Block]:
        head = Block(-1, 0)
        tail = head
        mode = False  # False == file, True == free space
        file_id = 0
        for num in map(int, self.getInput()):
            tail.right = Block(-1 if mode else file_id, num, tail)
            tail = tail.right

            if mode:
                file_id += 1
            mode = not mode

        head = head.right
        head.left = None
        return head, tail

    def part1(self) -> Any:
        head, tail = self.parse_input()

        index = head
        while index is not None:
            if index.file_id != -1:
                index = index.right
                continue

            if tail.file_id == -1:
                tail = tail.left
                tail.right = None

            if tail.file_length > index.file_length:
                index.file_id = tail.file_id
                tail.file_length -= index.file_length
            elif tail.file_length < index.file_length:
                space_left = index.file_length - tail.file_length
                index.file_id = tail.file_id
                index.file_length = tail.file_length
                tail = tail.left
                tail.right = None
                new_space = Block(-1, space_left, index)
                new_space.right = index.right
                index.right = new_space
            else:
                index.file_id = tail.file_id
                tail = tail.left
                tail.right = None

            index = index.right

        return get_checksum(head)

    def part2(self) -> Any:
        head, tail = self.parse_input()

        index = tail
        while index is not None:
            search = head
            while search != index:
                if search.file_id == -1 and search.file_length >= index.file_length:
                    new_free_space = Block(-1, index.file_length, index.left)
                    new_free_space.right = index.right
                    index.left.right = new_free_space
                    if index.right is not None:
                        index.right.left = new_free_space

                    search.file_id = index.file_id
                    if search.file_length > index.file_length:
                        space_left = search.file_length - index.file_length
                        search.file_length = index.file_length
                        new_space = Block(-1, space_left, search)
                        new_space.right = search.right
                        search.right.left = new_space
                        search.right = new_space

                    index = new_free_space
                    break
                search = search.right

            index = index.left

        while tail.file_id == -1:
            tail = tail.left
            tail.right = None

        return get_checksum(head)


if __name__ == "__main__":
    day = Day(2024, 9)
    day.run(verbose=True)