from typing import TypeVar, Generic, Iterable, Iterator, List
from .defs import Token, TokenType


T = TypeVar("T")
class PeekableIterator(Iterator[T]):
    _peeked: List[T]
    last_item: T | None

    def __init__(self, iterable: Iterable[T]) -> None:
        self.iter = iterable
        self._peeked = list()
        self.last_item = None

    def peek(self, offset: int = 0):
        while len(self._peeked) <= offset:
            try:
                self._peeked.append(next(self.iter))
            except StopIteration:
                return None
        
        return self._peeked[offset]

    def __next__(self) -> T:
        if len(self._peeked) > 0:
            item = self._peeked.pop(0)
        else:
            item = next(self.iter)
        self.last_item = item
        return item

    def __iter__(self) -> Iterator[T]:
        return self

    def next(self) -> T:
        try:
            return next(self)
        except StopIteration:
            return None

    def has_next(self):
        return self.peek() is not None


class ParserIterator(PeekableIterator[Token]):
    def __init__(self, iterable: Iterable[Token]) -> None:
        super().__init__(t for t in iterable if t.kind not in (TokenType.LineComment, TokenType.MultiComment))
        self.ignore_newline = False

    def peek(self, offset: int = 0):
        while len(self._peeked) <= offset:
            try:
                self._peeked.append(next(self.iter))
            except StopIteration:
                return None
        
        token = self._peeked[offset]

        if self.ignore_newline and token.kind == TokenType.EOL:
            return self.peek(offset=offset+1)

        return token

    def __next__(self) -> T:
        if len(self._peeked) > 0:
            item = self._peeked.pop(0)
        else:
            item = next(self.iter)
        self.last_item = item

        if self.ignore_newline and item.kind == TokenType.EOL:
            return next(self)

        return item