import struct
from skyhookdmclient.skyhook_common import *
import awkward

class SkyhookDM:
    def __init__(self):
        self.client = None
        self.addr = None

    def connect(self, ip, ceph_pool_name):
        addr = ip+':8786'
        self.addr = addr
        client = Client(addr)
        self.client = client
        self.ceph_pool = ceph_pool_name

    def saveDatabase(self, db):

        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res
        
        fu = self.client.submit(runOnDriver, db.serialize(), db.getName(), self.ceph_pool)
        result = fu.result()
        return result


    def write_full(self, table_name, table):
        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res

        # add the arrow table metadata , assuming it has at lease one col, need to handle exceptions
        meta_obj = TableMeta('0', '1', '2', SkyFormatType.SFT_ARROW, "Need to be extracted from the fields of the table (to do)", 'n/a', table_name, str(len(table.columns[0])))
        sche_meta = meta_obj.getTableMeta()
        schema = table.schema
        schema = schema.with_metadata(sche_meta)
            
        # Serialize arrow table to bytes
        batches = table.to_batches()
        sink = pa.BufferOutputStream()
        writer = pa.RecordBatchStreamWriter(sink, schema)
        for batch in batches:
            writer.write_batch(batch)   
        buff = sink.getvalue()
        buff_bytes = buff.to_pybytes()

        # add FB_Meta wrapper
        buff_bytes = addFB_Meta(buff_bytes)

        # write to Ceph
        fu = self.client.submit(runOnDriver, buff_bytes, table_name, self.ceph_pool)

        result = fu.result()
        if result is True:
            return True
        
        return False


    # Write the arrow table to Ceph
    def write(self, table, db_name, prefix, table_name, type= 'arrow', partitioner = None, mode = 'create|overwrite'):
        
        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res

        id_array = range(len(table.columns[0]))
        event_id_col = pa.field('EVENT_ID', pa.int64())

        object_names = []

        # Create a table for each column
        for i in range(len(table.columns)):
            field = table.field(0)
            schema = pa.schema([event_id_col, table.field(i)])
            sub_table = pa.Table.from_arrays([id_array, table.columns[i]],schema = schema)
            
            # Serialize arrow table to bytes
            batches = sub_table.to_batches()
            sink = pa.BufferOutputStream()
            writer = pa.RecordBatchStreamWriter(sink, sub_table.schema)

            for batch in batches:
                writer.write_batch(batch)
            
            buff = sink.getvalue()
            buff_bytes = buff.to_pybytes()
            buff_bytes = addFB_Meta(buff_bytes)

            sub_table_name = db_name + '#' + prefix + '#' + table_name + '#' + table.column_names[i] + '.0.0'
            fu = self.client.submit(runOnDriver, buff_bytes, sub_table_name, self.ceph_pool)
            result = fu.result()
            if result is True:
                object_names.append(sub_table_name)
        
        return object_names


    def writeDataset(self, path, dstname):
        def runOnDriver(path, dstname, poolname, addr ):
            # import skyhook_driver as sd
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeDataset(path, dstname, addr, poolname)
            return res

        fu = self.client.submit(runOnDriver, path, dstname, self.ceph_pool, self.addr)
        result = fu.result()
        return result

    def getDataset(self, name):
        def runOnDriver(name, ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.getDataset(name, ceph_pool)
            return res

        fu = self.client.submit(runOnDriver, name, self.ceph_pool)
        result = fu.result()
        
        data = json.loads(result)

        if 'type' in data.keys():
            return data

        else:
            files = [] 
            for item in data['files']:
                file = File(item['name'], item['file_attributes'], item['file_schema'], name, str(item['ROOTDirectory']))
                files.append(file)

            dataset = Dataset(data['dataset_name'], data['size'], files)
            return dataset

    def runQuery(self, obj, querystr):
        #limit just to 1 obj
        command_template = ['--wthreads', '1',  '--qdepth', '120',  '--query', 'hep', '--pool', self.ceph_pool, '--output-format', 'SFT_PYARROW_BINARY', '--num-objs', '1', '--subpartitions', '10']

        def generateQueryCommand(file, querystr):
            cmds = []
            prefix = file.dataset + '#' + str(file.name) + '#' + str(file.ROOTDirectory)
            brs = querystr.split('project')[-1].split(',')
            obj_num = 0

            for br in brs:
                br = br.strip()
                elems = br.split('.')
                br_name = elems[-1]  
                elems.remove(br_name)
                local_prefix = ''
                for elem in elems:
                    local_prefix = local_prefix + '#' + elem.strip()
                obj_prefix = prefix + local_prefix + '#'
                data_schema = ''

                f_schema = file.getSchema()
                found = False

                for i in range(len(elems)):
                    for j in range(len(f_schema['children'])):
                        ch_sche = f_schema['children'][j]
                        if elems[i].strip() == ch_sche['name']:
                            f_schema = ch_sche
                            break

                obj_num = len(f_schema['children'])

                for m in range(len(f_schema['children'])):
                    ch_sche = f_schema['children'][m]
                    if br_name.strip() == ch_sche['name']:
                        found = True
                        f_schema = ch_sche
                        #limit the obj num to 1
                        obj_num = m
                        break

                if found:
                    cmd = command_template.copy()
                    cmd.append('--start-obj')
                    cmd.append(str(obj_num))

                    data_schema = '0 4 0 0 EVENT_ID;' + f_schema['data_schema'] + ';'
                    cmd.append('--data-schema')
                    cmd.append(data_schema)

                    cmd.append('--project')
                    cmd.append('event_id,'+ br_name.strip())

                    cmd.append('--oid-prefix')
                    cmd.append(obj_prefix)
                    
                    cmd.append('--table-name')
                    cmd.append(br_name.strip())
                    
                    #limit the obj num to 1
                    cmds.append(cmd)
            return cmds



        def exeQuery(command):
            import subprocess

            prog = '/mnt/sdb/skyhookdm-ceph/build/bin/run-query'
            cmd = []
            cmd.append(prog)
            cmd.extend(command)
            
            p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
            result = p.communicate()[0]

            return result


        def _mergeTables(tables):
            bigtable = None
            for table in tables:
                if bigtable == None:
                    bigtable = table.remove_column(0)
                else:
                    bigtable = bigtable.append_column(table.field(1), table.columns[1])

            return bigtable.replace_schema_metadata(None)


        def fileQuery(obj, querystr):
            cmds = generateQueryCommand(obj, querystr)
            futures = []
            for command in cmds:
                future = self.client.submit(exeQuery, command)
                futures.append(future)

            tablestreams = self.client.gather(futures)
            tables = []

            for tablestream in tablestreams:
                sizebf = None
                stream_length = len(tablestream)
                cursor = 0
                batches = []

                while True:
                    if cursor == stream_length:
                        break
                    sizebf = tablestream[cursor:cursor+8]
                    cursor = cursor + 8
                    size = struct.unpack("<i",sizebf)[0]
                    stream = tablestream[cursor:cursor+size]
                    cursor = cursor + size
                    reader = pa.ipc.open_stream(stream)
                    for b in reader:
                        batches.append(b)

                #sort batches
                def mykey(batch):
                    tb = pa.Table.from_batches([batch.slice(0,1)])
                    return tb.columns[0][0].as_py()

                batches = sorted(batches,key=mykey)

                table = pa.Table.from_batches(batches)
                tables.append(table)

            res = _mergeTables(tables)

            res = LazyDataframe(res)
            return res


        if 'File' in str(obj):
            res = fileQuery(obj, querystr)
            return res

        if 'Dataset' in str(obj):

            rtfiles = obj.getFiles()
            global_tables = []
            futures_set = []

            for rtfile in rtfiles:
                cmds = generateQueryCommand(rtfile, querystr)
                futures = []
                for command in cmds:
                    future = self.client.submit(exeQuery, command)
                    futures.append(future)
                futures_set.append(futures)

            for futures in futures_set:

                tablestreams = self.client.gather(futures)
                tables = []

                for tablestream in tablestreams:
                    sizebf = None
                    stream_length = len(tablestream)
                    cursor = 0
                    batches = []

                    while True:
                        if cursor == stream_length:
                            break
                        sizebf = tablestream[cursor:cursor+8]
                        cursor = cursor + 8
                        size = struct.unpack("<i",sizebf)[0]
                        stream = tablestream[cursor:cursor+size]
                        cursor = cursor + size
                        reader = pa.ipc.open_stream(stream)
                        for b in reader:
                            batches.append(b)

                    #sort batches
                    def mykey(batch):
                        tb = pa.Table.from_batches([batch.slice(0,1)])
                        return tb.columns[0][0].as_py()

                    batches = sorted(batches,key=mykey)

                    table = pa.Table.from_batches(batches)
                    tables.append(table)

                res = _mergeTables(tables)
                res = LazyDataframe(res)
                global_tables.append(res)

            return global_tables

        return None


class LazyDataframe:

    def __init__(self, arr_table, entrysteps = -1):
        self._arr_table = arr_table
        self._entrysteps = entrysteps
        self._batches = None
        self._current_index = 0

    def __iter__(self):
        return self


    def __next__(self):
        if self._batches is None:
            raise Exception('Please set entrysteps using set_entrysteps() function before use!') 

        if self._current_index >= len(self._batches): 
            self._current_index = 0
            raise StopIteration

        sub_table = pa.Table.from_batches(self._batches[self._current_index:self._current_index + 1])
        self._current_index = self._current_index + 1

        chunk = sub_table.to_pydict()

        for item in chunk:
            chunk[item] = awkward.fromiter(chunk[item])

        return chunk

    def set_entrysteps(self, entrysteps):
        if self._batches is not None:
            self._arr_table = pa.Table.from_batches(self._batches)

        self._entrysteps = entrysteps
        self._batches = self._arr_table.to_batches(entrysteps)


class Database:
    def __init__(self, db_name, metadata = None):
        self._name = db_name
        self.schema = {}
        self.schema['db_name'] = db_name
        self.schema['metadata'] = metadata
        self.tables = []
    

    def createTable(self, prefix, table_name, table_schema, tag = '', metadata = None):
        schema = table_schema

        table = {}
        table['prefix'] = prefix
        table['table_name'] = table_name

        col_list = []

        for i in range(len(schema.names)):
            col = {}
            col['name'] = schema.field(i).name
            col['type'] = str(schema.field(i).type)
            col_list.append(col)
        
        table['schema'] = col_list
        table['metadata'] = metadata

        self.tables.append(table)

    def serialize(self):
        self.schema['tables'] = self.tables
        json_str = json.dumps(self.schema,indent=4)
        buff_bytes = bytes(json_str,'utf-8')
        return buff_bytes

    def getName(self):
        return self._name
import struct
from skyhookdmclient.skyhook_common import *
import awkward

class SkyhookDM:
    def __init__(self):
        self.client = None
        self.addr = None

    def connect(self, ip, ceph_pool_name):
        addr = ip+':8786'
        self.addr = addr
        client = Client(addr)
        self.client = client
        self.ceph_pool = ceph_pool_name

    def saveDatabase(self, db):

        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res
        
        fu = self.client.submit(runOnDriver, db.serialize(), db.getName(), self.ceph_pool)
        result = fu.result()
        return result


    def write_full(self, table_name, table):
        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res

        # add the arrow table metadata , assuming it has at lease one col, need to handle exceptions
        meta_obj = TableMeta('0', '1', '2', SkyFormatType.SFT_ARROW, "Need to be extracted from the fields of the table (to do)", 'n/a', table_name, str(len(table.columns[0])))
        sche_meta = meta_obj.getTableMeta()
        schema = table.schema
        schema = schema.with_metadata(sche_meta)
            
        # Serialize arrow table to bytes
        batches = table.to_batches()
        sink = pa.BufferOutputStream()
        writer = pa.RecordBatchStreamWriter(sink, schema)
        for batch in batches:
            writer.write_batch(batch)   
        buff = sink.getvalue()
        buff_bytes = buff.to_pybytes()

        # add FB_Meta wrapper
        buff_bytes = addFB_Meta(buff_bytes)

        # write to Ceph
        fu = self.client.submit(runOnDriver, buff_bytes, table_name, self.ceph_pool)

        result = fu.result()
        if result is True:
            return True
        
        return False


    # Write the arrow table to Ceph
    def write(self, table, db_name, prefix, table_name, type= 'arrow', partitioner = None, mode = 'create|overwrite'):
        
        def runOnDriver(buff_bytes, name,  ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeToCeph(buff_bytes, name, ceph_pool)
            return res

        id_array = range(len(table.columns[0]))
        event_id_col = pa.field('EVENT_ID', pa.int64())

        object_names = []

        # Create a table for each column
        for i in range(len(table.columns)):
            field = table.field(0)
            schema = pa.schema([event_id_col, table.field(i)])
            sub_table = pa.Table.from_arrays([id_array, table.columns[i]],schema = schema)
            
            # Serialize arrow table to bytes
            batches = sub_table.to_batches()
            sink = pa.BufferOutputStream()
            writer = pa.RecordBatchStreamWriter(sink, sub_table.schema)

            for batch in batches:
                writer.write_batch(batch)
            
            buff = sink.getvalue()
            buff_bytes = buff.to_pybytes()
            buff_bytes = addFB_Meta(buff_bytes)

            sub_table_name = db_name + '#' + prefix + '#' + table_name + '#' + table.column_names[i] + '.0.0'
            fu = self.client.submit(runOnDriver, buff_bytes, sub_table_name, self.ceph_pool)
            result = fu.result()
            if result is True:
                object_names.append(sub_table_name)
        
        return object_names


    def writeDataset(self, path, dstname):
        def runOnDriver(path, dstname, poolname, addr ):
            # import skyhook_driver as sd
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.writeDataset(path, dstname, addr, poolname)
            return res

        fu = self.client.submit(runOnDriver, path, dstname, self.ceph_pool, self.addr)
        result = fu.result()
        return result

    def getDataset(self, name):
        def runOnDriver(name, ceph_pool):
            from skyhookdmdriver import skyhook_driver as sd
            res = sd.getDataset(name, ceph_pool)
            return res

        fu = self.client.submit(runOnDriver, name, self.ceph_pool)
        result = fu.result()
        
        data = json.loads(result)

        if 'type' in data.keys():
            return data

        else:
            files = [] 
            for item in data['files']:
                file = File(item['name'], item['file_attributes'], item['file_schema'], name, str(item['ROOTDirectory']))
                files.append(file)

            dataset = Dataset(data['dataset_name'], data['size'], files)
            return dataset

    def runQuery(self, obj, querystr):
        #limit just to 1 obj
        command_template = ['--wthreads', '1',  '--qdepth', '120',  '--query', 'hep', '--pool', self.ceph_pool, '--output-format', 'SFT_PYARROW_BINARY', '--num-objs', '1', '--subpartitions', '10']

        def generateQueryCommand(file, querystr):
            cmds = []
            prefix = file.dataset + '#' + str(file.name) + '#' + str(file.ROOTDirectory)
            brs = querystr.split('project')[-1].split(',')
            obj_num = 0

            for br in brs:
                br = br.strip()
                elems = br.split('.')
                br_name = elems[-1]  
                elems.remove(br_name)
                local_prefix = ''
                for elem in elems:
                    local_prefix = local_prefix + '#' + elem.strip()
                obj_prefix = prefix + local_prefix + '#'
                data_schema = ''

                f_schema = file.getSchema()
                found = False

                for i in range(len(elems)):
                    for j in range(len(f_schema['children'])):
                        ch_sche = f_schema['children'][j]
                        if elems[i].strip() == ch_sche['name']:
                            f_schema = ch_sche
                            break

                obj_num = len(f_schema['children'])

                for m in range(len(f_schema['children'])):
                    ch_sche = f_schema['children'][m]
                    if br_name.strip() == ch_sche['name']:
                        found = True
                        f_schema = ch_sche
                        #limit the obj num to 1
                        obj_num = m
                        break

                if found:
                    cmd = command_template.copy()
                    cmd.append('--start-obj')
                    cmd.append(str(obj_num))

                    data_schema = '0 4 0 0 EVENT_ID;' + f_schema['data_schema'] + ';'
                    cmd.append('--data-schema')
                    cmd.append(data_schema)

                    cmd.append('--project')
                    cmd.append('event_id,'+ br_name.strip())

                    cmd.append('--oid-prefix')
                    cmd.append(obj_prefix)
                    
                    cmd.append('--table-name')
                    cmd.append(br_name.strip())
                    
                    #limit the obj num to 1
                    cmds.append(cmd)
            return cmds



        def exeQuery(command):
            import subprocess

            prog = '/mnt/sdb/skyhookdm-ceph/build/bin/run-query'
            cmd = []
            cmd.append(prog)
            cmd.extend(command)
            
            p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
            result = p.communicate()[0]

            return result


        def _mergeTables(tables):
            bigtable = None
            for table in tables:
                if bigtable == None:
                    bigtable = table.remove_column(0)
                else:
                    bigtable = bigtable.append_column(table.field(1), table.columns[1])

            return bigtable.replace_schema_metadata(None)


        def fileQuery(obj, querystr):
            cmds = generateQueryCommand(obj, querystr)
            futures = []
            for command in cmds:
                future = self.client.submit(exeQuery, command)
                futures.append(future)

            tablestreams = self.client.gather(futures)
            tables = []

            for tablestream in tablestreams:
                sizebf = None
                stream_length = len(tablestream)
                cursor = 0
                batches = []

                while True:
                    if cursor == stream_length:
                        break
                    sizebf = tablestream[cursor:cursor+8]
                    cursor = cursor + 8
                    size = struct.unpack("<i",sizebf)[0]
                    stream = tablestream[cursor:cursor+size]
                    cursor = cursor + size
                    reader = pa.ipc.open_stream(stream)
                    for b in reader:
                        batches.append(b)

                #sort batches
                def mykey(batch):
                    tb = pa.Table.from_batches([batch.slice(0,1)])
                    return tb.columns[0][0].as_py()

                batches = sorted(batches,key=mykey)

                table = pa.Table.from_batches(batches)
                tables.append(table)

            res = _mergeTables(tables)

            res = LazyDataframe(res)
            return res


        if 'File' in str(obj):
            res = fileQuery(obj, querystr)
            return res

        if 'Dataset' in str(obj):

            rtfiles = obj.getFiles()
            global_tables = []
            futures_set = []

            for rtfile in rtfiles:
                cmds = generateQueryCommand(rtfile, querystr)
                futures = []
                for command in cmds:
                    future = self.client.submit(exeQuery, command)
                    futures.append(future)
                futures_set.append(futures)

            for futures in futures_set:

                tablestreams = self.client.gather(futures)
                tables = []

                for tablestream in tablestreams:
                    sizebf = None
                    stream_length = len(tablestream)
                    cursor = 0
                    batches = []

                    while True:
                        if cursor == stream_length:
                            break
                        sizebf = tablestream[cursor:cursor+8]
                        cursor = cursor + 8
                        size = struct.unpack("<i",sizebf)[0]
                        stream = tablestream[cursor:cursor+size]
                        cursor = cursor + size
                        reader = pa.ipc.open_stream(stream)
                        for b in reader:
                            batches.append(b)

                    #sort batches
                    def mykey(batch):
                        tb = pa.Table.from_batches([batch.slice(0,1)])
                        return tb.columns[0][0].as_py()

                    batches = sorted(batches,key=mykey)

                    table = pa.Table.from_batches(batches)
                    tables.append(table)

                res = _mergeTables(tables)
                res = LazyDataframe(res)
                global_tables.append(res)

            return global_tables

        return None


class LazyDataframe:

    def __init__(self, arr_table, entrysteps = -1):
        self._arr_table = arr_table
        self._entrysteps = entrysteps
        self._batches = None
        self._current_index = 0

    def __iter__(self):
        return self


    def __next__(self):
        if self._batches is None:
            raise Exception('Please set entrysteps using set_entrysteps() function before use!') 

        if self._current_index >= len(self._batches): 
            self._current_index = 0
            raise StopIteration

        sub_table = pa.Table.from_batches(self._batches[self._current_index:self._current_index + 1])
        self._current_index = self._current_index + 1

        chunk = sub_table.to_pydict()

        for item in chunk:
            chunk[item] = awkward.fromiter(chunk[item])

        return chunk

    def set_entrysteps(self, entrysteps):
        if self._batches is not None:
            self._arr_table = pa.Table.from_batches(self._batches)

        self._entrysteps = entrysteps
        self._batches = self._arr_table.to_batches(entrysteps)


class Database:
    def __init__(self, db_name, metadata = None):
        self._name = db_name
        self.schema = {}
        self.schema['db_name'] = db_name
        self.schema['metadata'] = metadata
        self.tables = []
    

    def createTable(self, prefix, table_name, table_schema, tag = '', metadata = None):
        schema = table_schema

        table = {}
        table['prefix'] = prefix
        table['table_name'] = table_name

        col_list = []

        for i in range(len(schema.names)):
            col = {}
            col['name'] = schema.field(i).name
            col['type'] = str(schema.field(i).type)
            col_list.append(col)
        
        table['schema'] = col_list
        table['metadata'] = metadata

        self.tables.append(table)

    def serialize(self):
        self.schema['tables'] = self.tables
        json_str = json.dumps(self.schema,indent=4)
        buff_bytes = bytes(json_str,'utf-8')
        return buff_bytes

    def getName(self):
        return self._name
# add the arrow table metadata , assuming it has at lease one col, need to handle exceptions # Serialize arrow table to bytes # add FB_Meta wrapper # write to Ceph # Write the arrow table to Ceph # Create a table for each column # Serialize arrow table to bytes # import skyhook_driver as sd #limit just to 1 obj #limit the obj num to 1 #limit the obj num to 1 #sort batches #sort batches