import math
from itertools import count, pairwise
from typing import Iterable, NamedTuple, override
from unittest import TestCase

from more_itertools import chunked, ilen

from puzzles._solver import Solver


class LinkedList[T]:
    __head: "LinkedListNode[T] | None"
    __tail: "LinkedListNode[T] | None"

    __slots__ = ("__head", "__tail")

    @property
    def head(self) -> "LinkedListNode[T] | None":
        if not self.__head and self.__tail:
            self.__head = self.__tail
        while self.__head and self.__head.left:
            self.__head = self.__head.left
        return self.__head

    @property
    def tail(self) -> "LinkedListNode[T] | None":
        if not self.__tail and self.__head:
            self.__tail = self.__head
        while self.__tail and self.__tail.right:
            self.__tail = self.__tail.right
        return self.__tail

    def __init__(self, *values: Iterable[T]) -> None:
        self.__head = None
        self.__tail = None

        for value in values:
            self.append(value)

    def append(self, value: T) -> None:
        if tail := self.tail:
            self.__tail = LinkedListNode(value, tail, None)
            tail.right = self.__tail
        else:
            self.__tail = LinkedListNode(value, None, None)

    def prepend(self, value: T) -> None:
        if head := self.head:
            self.__head = LinkedListNode(value, None, head)
            head.left = self.__head
        else:
            self.__head = LinkedListNode(value, None, None)

    def pop_tail(self) -> T | None:
        if (tail := self.tail) is None:
            raise IndexError("pop from empty LinkedList")

        if tail.left:
            tail.left.right = None

        self.__tail = tail.left
        tail.left = None

        return tail.value

    def pop_head(self) -> T:
        if (head := self.head) is None:
            raise IndexError("pop from empty LinkedList")

        if head.right:
            head.right.left = None

        self.__head = head.right
        head.right = None

        return head.value

    def copy(self) -> "LinkedList[T]":
        return LinkedList(*self)

    def __iter__(self) -> Iterable[T]:
        node = self.head
        while node is not None:
            yield node.value
            node = node.right

    def __reversed__(self) -> Iterable[T]:
        node = self.tail
        while node is not None:
            yield node.value
            node = node.left

    def __len__(self) -> int:
        return ilen(self)

    def __str__(self) -> str:
        return f"LinkedList({', '.join(map(str, self))})"


class LinkedListNode[T]:
    value: T
    left: "LinkedListNode[T] | None"
    right: "LinkedListNode[T] | None"

    __slots__ = ("value", "left", "right")

    def __init__(
        self,
        value: T,
        left: "LinkedListNode[T] | None",
        right: "LinkedListNode[T] | None",
    ) -> None:
        self.value = value
        self.left = left
        self.right = right

    def insert_left(self, value: T) -> None:
        node = LinkedListNode(value, self.left, self)
        if self.left:
            self.left.right = node
        self.left = node

    def insert_right(self, value: T) -> None:
        node = LinkedListNode(value, self, self.right)
        if self.right:
            self.right.left = node
        self.right = node

    def pop(self) -> T:
        if self.left:
            self.left.right = self.right
        if self.right:
            self.right.left = self.left
        return self.value


class TestLinkedList(TestCase):
    def test(self):
        linked_list = LinkedList(1, 2, 3, 4, 5)
        self.assertEqual(linked_list.head.value, 1)
        self.assertEqual(linked_list.tail.value, 5)
        self.assertListEqual(list(linked_list), [1, 2, 3, 4, 5])
        self.assertListEqual(list(reversed(linked_list)), [5, 4, 3, 2, 1])
        self.assertEqual(len(linked_list), 5)

        self.assertEqual(linked_list.pop_tail(), 5)
        self.assertEqual(linked_list.tail.value, 4)
        self.assertEqual(linked_list.pop_head(), 1)
        self.assertEqual(linked_list.head.value, 2)
        linked_list.append(6)
        self.assertEqual(linked_list.tail.value, 6)
        linked_list.prepend(7)
        self.assertEqual(linked_list.head.value, 7)
        self.assertListEqual(list(linked_list), [7, 2, 3, 4, 6])

        middle = linked_list.tail.left.left
        self.assertEqual(middle.value, 3)
        middle.insert_left(8)
        self.assertListEqual(list(linked_list), [7, 2, 8, 3, 4, 6])
        middle.insert_right(9)
        self.assertListEqual(list(linked_list), [7, 2, 8, 3, 9, 4, 6])
        self.assertEqual(middle.pop(), 3)
        self.assertListEqual(list(linked_list), [7, 2, 8, 9, 4, 6])


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


def compact_drive(drive: LinkedList[Block]) -> LinkedList[Block]:
    if not len(drive):
        return drive

    drive = drive.copy()
    moving = drive.pop_tail()
    cursor = drive.head
    assert cursor is not None

    while cursor.right is not None:
        a = cursor.value
        b = cursor.right.value
        free_space = b.position - a.position - a.size

        if not free_space:
            cursor = cursor.right
            continue

        insertion = Block(
            moving.id,
            a.position + a.size,
            min(free_space, moving.size),
        )

        cursor.insert_right(insertion)
        cursor = cursor.right

        if insertion.size < moving.size:
            moving = Block(
                moving.id,
                moving.position,
                moving.size - insertion.size,
            )
        else:
            moving = drive.pop_tail()

    cursor.insert_right(
        Block(
            moving.id,
            cursor.value.position + cursor.value.size,
            moving.size,
        )
    )

    return drive


def compact_drive_without_fragmentation(drive: LinkedList[Block]) -> LinkedList[Block]:
    if not len(drive):
        return drive

    drive = drive.copy()
    moving = drive.tail

    while moving:
        cursor = drive.head

        while cursor != moving:
            a = cursor.value
            b = cursor.right.value
            free_space = b.position - a.position - a.size

            if free_space >= moving.value.size:
                moving.pop()
                cursor.insert_right(
                    Block(
                        moving.value.id,
                        a.position + a.size,
                        moving.value.size,
                    )
                )
                break

            cursor = cursor.right

        moving = moving.left

    return drive


def hash_drive(drive: LinkedList[Block]) -> int:
    return sum(
        block.id * block.size * (block.position * 2 + block.size - 1) // 2
        for block in drive
    )


def print_drive(drive: LinkedList[Block]) -> None:
    snapshot = []
    for a, b in pairwise(drive):
        snapshot.append(str(a.id) * a.size)
        snapshot.append("." * (b.position - a.position - a.size))
    snapshot.append(str(drive.tail.value.id) * drive.tail.value.size)
    print("".join(snapshot))


class DayNineSolver(Solver):
    drive: LinkedList[Block]

    @override
    def __init__(self, puzzle_input: str):
        self.drive = LinkedList()
        snapshot = map(int, puzzle_input.strip())
        cursor_position = 0
        current_id = count()

        for pair in chunked(snapshot, 2):
            if data_size := pair[0]:
                id = next(current_id)
                self.drive.append(Block(id, cursor_position, data_size))
                cursor_position += data_size
            if len(pair) == 2:
                cursor_position += pair[1]

    @override
    def solve_p1(self) -> int:
        compacted_drive = compact_drive(self.drive)
        return hash_drive(compacted_drive)

    @override
    def solve_p2(self) -> int:
        compacted_drive = compact_drive_without_fragmentation(self.drive)
        return hash_drive(compacted_drive)


class TestDayNineSolver(TestCase):
    def test(self):
        solver = DayNineSolver("2333133121414131402\n")
        self.assertEqual(solver.solve_p1(), 1928)
        self.assertEqual(solver.solve_p2(), 2858)