from typing import Iterator, Optional
from queryservice_client.protobuf.rdbms.query_pb2 import QueryResponse, QueryResultObject
from queryservice_client.stream.immutablelist import ImmutableList
from queryservice_client.stream.reconnect import Reconnect
from queryservice_client.stream.cancelled import Cancelled

class Stream:
    def __init__(self, stream: Iterator[QueryResponse]):
        self.stream = stream
        self.buffer = ImmutableList()
        self._peek: Optional[QueryResultObject] = None

    def has_next(self) -> bool:
        is_next = self._peek is not None or not self.buffer.is_empty()
        if not is_next:
            try:
                next(self.stream)
                is_next = True
            except StopIteration:
                return False
        return is_next

    def __next__(self) -> QueryResultObject:
        if self._peek is not None:
            next_obj = self._peek
            self._peek = None
            return next_obj
        elif not self.buffer.is_empty():
            return self.buffer.pop()
        else:
            while True:
                result = next(self.stream)
                if result.cancelled:
                    raise Cancelled()
                if result.restart:
                    raise Reconnect()
                objects = result.queryResult.objects
                if not objects:
                    continue
                self.buffer.set(objects)
                return self.buffer.pop()

    def peek(self) -> QueryResultObject:
        if self._peek is not None:
            return self._peek
        try:
            self._peek = next(self)
        except StopIteration:
            return None
        return self._peek