from typing import Iterator
from queryservice_client.stream.segments.metadata import Metadata
from queryservice_client.stream.resultsetstream import ResultSetStream
from queryservice_client.stream.segments.row import Row
from queryservice_client.stream.rowstream import RowStream
from queryservice_client.stream.types.all_data_type import Data
from queryservice_client.stream.metadatastream import MetadataStream

class ResultSet(Iterator[Row]):
    def __init__(self, stream: ResultSetStream):
        self.stream = stream
        self.update_count = -1
        self.metadata = Metadata()
        self.eager_init()

    def eager_init(self):
        md = MetadataStream(self.stream)
        for metadata in md:
            if metadata.projection is not None:
                self.metadata.add(metadata.projection)
            else:
                self.update_count = metadata.updateCount

    def get_metadata(self):
        return self.metadata

    def get_update_count(self):
        return self.update_count

    def __iter__(self):
        return self

    def __next__(self):
        stream = False
        row = Row()
        for data in RowStream(self.stream):
            row.append(Data.create(data))
            stream = True
        if not stream:
            raise StopIteration
        return row

    def has_next(self):
        return self.stream.has_next()