import collections
import csv
import io
import json

import splunklib.searchcommands.internals
from splunklib import six


class Chunk(object):
    def __init__(self, version, meta, data):
        self.version = six.ensure_str(version)
        self.meta = json.loads(meta)
        self.data = csv.DictReader(io.StringIO(data.decode("utf-8")), dialect=splunklib.searchcommands.internals.CsvDialect)


class ChunkedDataStreamIter(collections.Iterator):
    def __init__(self, chunk_stream):
        self.chunk_stream = chunk_stream

    def __next__(self):
        return self.next()

    def next(self):
        try:
            return self.chunk_stream.read_chunk()
        except EOFError:
            raise StopIteration


class ChunkedDataStream(collections.Iterable):
    def __iter__(self):
        return ChunkedDataStreamIter(self)

    def __init__(self, stream):
        empty = stream.read(0)
        assert isinstance(empty, bytes)
        self.stream = stream

    def read_chunk(self):
        header = self.stream.readline()
        if len(header) == 0:
            raise EOFError

        while len(header) > 0 and header.strip() == b'':
            header = self.stream.readline()  # Skip empty lines
        version, meta, data = header.rstrip().split(b',')
        metabytes = self.stream.read(int(meta))
        databytes = self.stream.read(int(data))
        print("sent")
        return Chunk(version, metabytes, databytes)


def build_chunk(kv, data=None):
    metadata = six.ensure_binary(json.dumps(kv), 'utf-8')
    data_output = _build_data_csv(data)
    return b"chunked 1.0,%d,%d\n%s%s" % (len(metadata), len(data_output), metadata, data_output)


def build_empty_searchinfo():
    return {
        'earliest_time': 0,
        'latest_time': 0,
        'search': "",
        'dispatch_dir': "",
        'sid': "",
        'args': [],
        'splunk_version': "42.3.4",
    }


def build_getinfo_chunk():
    return build_chunk({'action': 'getinfo', 'preview': False, 'searchinfo': build_empty_searchinfo()})


def build_data_chunk(data, finished=True):
    return build_chunk({'action': 'execute', 'finished': finished}, data)


def _build_data_csv(data):
    if data is None:
        return b''
    if isinstance(data, bytes):
        return data
    if six.PY2:
        csvout = io.BytesIO()
    else:
        csvout = io.StringIO()

    headers = set()
    for datum in data:
        if six.PY2:
            headers.update(datum.keys())
        else:
            headers.update(datum.keys())
    writer = csv.DictWriter(csvout, headers, dialect=splunklib.searchcommands.internals.CsvDialect)
    writer.writeheader()
    for datum in data:
        if six.PY2:
            writer.writerow(datum)
        else:
            writer.writerow(datum)
    return six.ensure_binary(csvout.getvalue())

@@ -1035,64 +1035,8 @@
 ne()
-%0A        if len(header) == 0:%0A            raise EOFError
 %0A%0A  
@@ -1153,16 +1153,73 @@
 ty lines
+%0A%0A        if len(header) == 0:%0A            raise EOFError
 %0A       
@@ -2279,214 +2279,86 @@
     
-if six.PY2:%0A        csvout = io.BytesIO()%0A    else:%0A        csvout = io.StringIO()%0A%0A    headers = set()%0A    for datum in data:%0A        if six.PY2:%0A            headers.update(datum.keys())%0A        else:%0A    
+csvout = splunklib.six.StringIO()%0A%0A    headers = set()%0A    for datum in data:%0A
     
@@ -2542,81 +2542,8 @@
 ta:%0A
-        if six.PY2:%0A            writer.writerow(datum)%0A        else:%0A    
     

address code review comments