"""VirtualAllocator: device-wide VA space management with free-list.

Allocations are page-aligned. Freed ranges are coalesced with adjacent
free blocks to allow larger re-allocations.
"""
from __future__ import annotations

import bisect


class VaAllocationError(Exception):
    pass


class VirtualAllocator:
    """Manages a contiguous VA address space with page-aligned alloc/free.

    Args:
        va_base: Start of the VA range.
        va_size: Total size of the VA range in bytes.
        page_size: Page granularity in bytes.
    """

    def __init__(
        self,
        va_base: int,
        va_size: int,
        page_size: int = 2 * 1024 * 1024,
    ) -> None:
        self._va_base = va_base
        self._va_size = va_size
        self._page_size = page_size
        self._used = 0
        # Free list: sorted list of (start, size) tuples
        self._free: list[tuple[int, int]] = [(va_base, va_size)]

    def _align_up(self, nbytes: int) -> int:
        """Round up to page boundary."""
        return ((nbytes + self._page_size - 1) // self._page_size) * self._page_size

    def alloc(self, nbytes: int) -> int:
        """Allocate a contiguous VA range. Returns the start VA."""
        aligned = self._align_up(nbytes)
        for i, (start, size) in enumerate(self._free):
            if size >= aligned:
                # Take from the beginning of this free block
                if size == aligned:
                    self._free.pop(i)
                else:
                    self._free[i] = (start + aligned, size - aligned)
                self._used += aligned
                return start
        raise VaAllocationError(
            f"Out of VA space: need {aligned}, largest free block "
            f"{max((s for _, s in self._free), default=0)}"
        )

    def free(self, va: int, nbytes: int) -> None:
        """Free a VA range and coalesce with adjacent free blocks."""
        aligned = self._align_up(nbytes)
        self._used -= aligned

        # Insert into sorted free list and coalesce
        new_start = va
        new_end = va + aligned

        # Find insertion point
        idx = bisect.bisect_left(self._free, (va,))

        # Try coalesce with previous block
        if idx > 0:
            prev_start, prev_size = self._free[idx - 1]
            if prev_start + prev_size == new_start:
                new_start = prev_start
                idx -= 1
                self._free.pop(idx)

        # Try coalesce with next block
        if idx < len(self._free):
            next_start, next_size = self._free[idx]
            if new_end == next_start:
                new_end = next_start + next_size
                self._free.pop(idx)

        self._free.insert(idx, (new_start, new_end - new_start))

    @property
    def used(self) -> int:
        return self._used

    @property
    def total(self) -> int:
        return self._va_size
