from dataclasses import dataclass
from typing import Callable, Iterable, Any

@dataclass
class Window[T]:
    """Window for AI processing."""

    past: list[T]
    """Part of data from previous window (can be empty)"""

    present: list[T]
    """Data for current window (can NOT be empty)"""

    future: list[T]
    """Part of data from the next window (can be empty)"""

    context: dict[str, Any]
    """Context that is preserved between LLM iterations"""


class Windowizer[T]:
    """This class builds windows from input data"""

    @staticmethod
    def _build_list_with_limits(items: Iterable[T], min_count: int, max_size: int, size_func: Callable[[T], int]) -> list[T]:
        result: list[T] = []
        total_size: int = 0
        for i in items:
            result.append(i)
            total_size += size_func(i)
            if len(result) < min_count:
                continue
            if total_size >= max_size:
                break
        return result

    @staticmethod
    def default_size_func(elem: T) -> int:
        """Default size function just returns length of the `repr` result."""
        return len(repr(elem))

    def __init__(self,
                 *,
                 main_min_count: int = 2,
                 main_max_size: int = 1024,
                 side_min_count: int = 2,
                 side_max_size: int = 256,
                 size_func: Callable[[T], int] = default_size_func) -> None:
        """Create Windowizer.
        
        Args:
            - main_min_count - minimum count of items in `present`
            - main_max_size - maximum total size of items in `present`
            - side_min_count - minimum count of items in `past` and `future`
            - side_max_size - maximum total size of items in `past` and `future`
            - size_func - function that will be used to get item size
        """
        self._main_min_count = main_min_count
        self._main_max_size = main_max_size
        self._side_min_count = side_min_count
        self._side_max_size = side_max_size
        self._size_func = size_func

    def windowize(self, items: list[T]) -> list[Window[T]]:
        """Builds windows from items.

        `min_count_*` has more priority than `max_size_*`. It would be possible
        to build windows of 0 items otherwise.
        """
        result: list[Window[T]] = []
        # create windows
        window_start = 0
        while window_start < len(items):
            window = Window[T](past=[], present=[], future=[], context={})
            # build `present`
            window.present = self._build_list_with_limits(
                items[window_start:],
                min_count=self._main_min_count,
                max_size=self._main_max_size,
                size_func=self._size_func
            )
            # build `past`
            window.past = self._build_list_with_limits(
                reversed(items[:window_start]),
                min_count=self._side_min_count,
                max_size=self._side_max_size,
                size_func=self._size_func
            )
            window.past.reverse()
            # build `future`
            window.future = self._build_list_with_limits(
                items[window_start+len(window.present):],
                min_count=self._side_min_count,
                max_size=self._side_max_size,
                size_func=self._size_func
            )
            window_start += len(window.present)
            result.append(window)
        # resulting windows
        return result