# General Writer to drive VDB General Loader

from sys import stdout, stderr
from struct import Struct, calcsize, pack

_headerFmt = Struct("8s 4I")
_simpleEventFmt = Struct("I")
_dataEventFmt = Struct("2I")
_writeStdOut = stdout.buffer.write
_makeDataEvent = _dataEventFmt.pack


def _padding(length: int) -> str:
    return f" {4 - (length & 0x03)}x" if length & 0x03 > 0 else ''


def _paddedFormat(fmt):
    return fmt + _padding(calcsize(fmt))


def _makeHeader():
    return _headerFmt.pack("NCBIgnld".encode('ascii'),
                           1, 2, _headerFmt.size, 0)


def _makeSimpleEvent(eid):
    """ used for { evt_end_stream, evt_open_stream, evt_next_row } """
    return _simpleEventFmt.pack(eid)


def _make1StringEvent(eid, str1):
    """ used for { evt_errmsg, evt_remote_path, evt_new_table, evt_software_name } """
    len1 = len(str1)
    fmt = _paddedFormat(f"I 1I {len1}s")
    return pack(fmt, eid, len1, str1)


def _make2StringEvent(eid, str1, str2):
    """ used for { evt_use_schema, evt_metadata_node } """
    len1 = len(str1)
    len2 = len(str2)
    fmt = _paddedFormat(f"I 2I {len1}s {len2}s")
    return pack(fmt, eid, len1, len2, str1, str2)

def _make3StringEvent(eid, str1, str2, str3):
    """ used for { evt_db_metadata_node_attr, evt_tbl_metadata_node_attr, evt_col_metadata_node_attr } """
    len1 = len(str1)
    len2 = len(str2)
    len3 = len(str3)
    fmt = _paddedFormat(f"I 3I {len1}s {len2}s {len3}s")
    return pack(fmt, eid, len1, len2, len3, str1, str2, str3)


def _makeColumnEvent(colid, tblid, bits, name):
    """ used for { evt_new_column } """
    lenn = len(name)
    fmt = _paddedFormat(f"I 3I {lenn}s")
    return pack(fmt, colid, tblid, bits, lenn, name)


class GeneralWriter:
    # pylint: disable=too-few-public-methods

    evt_bad_event       = 0

    evt_errmsg          = (1 << 24) + evt_bad_event
    evt_end_stream      = (1 << 24) + evt_errmsg

    evt_remote_path     = (1 << 24) + evt_end_stream
    evt_use_schema      = (1 << 24) + evt_remote_path
    evt_new_table       = (1 << 24) + evt_use_schema
    evt_new_column      = (1 << 24) + evt_new_table
    evt_open_stream     = (1 << 24) + evt_new_column

    evt_cell_default    = (1 << 24) + evt_open_stream
    evt_cell_data       = (1 << 24) + evt_cell_default
    evt_next_row        = (1 << 24) + evt_cell_data

    evt_move_ahead      = (1 << 24) + evt_next_row      # this one is not used here
    evt_errmsg2         = (1 << 24) + evt_move_ahead    # this one is not used here
    evt_remote_path2    = (1 << 24) + evt_errmsg2       # this one is not used here
    evt_use_schema2     = (1 << 24) + evt_remote_path2  # this one is not used here
    evt_new_table2      = (1 << 24) + evt_use_schema2   # this one is not used here
    evt_cell_default2   = (1 << 24) + evt_new_table2    # this one is not used here
    evt_cell_data2      = (1 << 24) + evt_cell_default2 # this one is not used here
    evt_empty_default   = (1 << 24) + evt_cell_data2    # this one is not used here

    # BEGIN VERSION 2 MESSAGES
    evt_software_name      = (1 << 24) + evt_empty_default
    evt_db_metadata_node   = (1 << 24) + evt_software_name
    evt_tbl_metadata_node  = (1 << 24) + evt_db_metadata_node
    evt_col_metadata_node  = (1 << 24) + evt_tbl_metadata_node
    evt_db_metadata_node2  = (1 << 24) + evt_col_metadata_node
    evt_tbl_metadata_node2 = (1 << 24) + evt_db_metadata_node2
    evt_col_metadata_node2 = (1 << 24) + evt_tbl_metadata_node2

    evt_add_mbr_db         = (1 << 24) + evt_col_metadata_node2
    evt_add_mbr_tbl        = (1 << 24) + evt_add_mbr_db

    evt_logmsg             = (1 << 24) + evt_add_mbr_tbl
    evt_progmsg            = (1 << 24) + evt_logmsg

    # BEGIN VERSION 3 MESSAGES
    evt_db_metadata_node_attr   = (1 << 24) + evt_progmsg
    evt_tbl_metadata_node_attr  = (1 << 24) + evt_db_metadata_node_attr
    evt_col_metadata_node_attr  = (1 << 24) + evt_tbl_metadata_node_attr
    evt_db_metadata_node_attr2  = (1 << 24) + evt_col_metadata_node_attr
    evt_tbl_metadata_node_attr2 = (1 << 24) + evt_db_metadata_node_attr2
    evt_col_metadata_node_attr2 = (1 << 24) + evt_tbl_metadata_node_attr2

    def errorMessage(self, message):
        _writeStdOut(_make1StringEvent(self.evt_errmsg, message.encode('utf-8')))

    def logMessage(self, message):
        _writeStdOut(_make1StringEvent(self.evt_logmsg, message.encode('utf-8')))

    def write(self, spec):
        tableId = spec['_tableId']
        for k in spec:
            if k[0] == '_':
                continue

            c = spec[k]
            if 'data' in c:
                data = c['data']
                if callable(data):
                    data = data()

                try:
                    self._writeColumnData(c['_columnId'], len(data), data)
                except Exception as e:
                    stderr.write(f"failed to write column #{c['_columnId']} {k}: {e}\n")
                    raise e
        self._writeNextRow(tableId)

    @classmethod
    def _writeHeader(cls, remoteDb, schemaFileName, schemaDbSpec):
        _writeStdOut(_makeHeader())
        _writeStdOut(_make1StringEvent(cls.evt_remote_path, remoteDb))
        _writeStdOut(_make2StringEvent(cls.evt_use_schema, schemaFileName, schemaDbSpec))

    @classmethod
    def _writeSoftwareName(cls, name, version): # name is any string, version is like "2.1.5"
        _writeStdOut(_make2StringEvent(cls.evt_software_name, name, version))

    @classmethod
    def _writeEndStream(cls):
        _writeStdOut(_makeSimpleEvent(cls.evt_end_stream))

    @classmethod
    def _writeNewTable(cls, tableId, table):
        _writeStdOut(_make1StringEvent(cls.evt_new_table + tableId, table))

    @classmethod
    def _writeNewColumn(cls, columnId, tableId, bits, column):
        _writeStdOut(_makeColumnEvent(cls.evt_new_column + columnId, tableId, bits, column))

    @classmethod
    def _writeOpenStream(cls):
        _writeStdOut(_makeSimpleEvent(cls.evt_open_stream))

    @classmethod
    def _writeColumnDefault(cls, colId, count, data):
        l = _writeStdOut(_makeDataEvent(cls.evt_cell_default + colId, count))
        l += _writeStdOut(data)
        _writeStdOut(pack(_padding(l)))

    @classmethod
    def _writeDbMetadata(cls, dbId, nodeName, nodeValue):
        _writeStdOut(_make2StringEvent(cls.evt_db_metadata_node + dbId, nodeName, nodeValue))

    @classmethod
    def _writeTableMetadata(cls, tblId, nodeName, nodeValue):
        _writeStdOut(_make2StringEvent(cls.evt_tbl_metadata_node + tblId, nodeName, nodeValue))

    @classmethod
    def _writeColumnMetadata(cls, colId, nodeName, nodeValue):
        _writeStdOut(_make2StringEvent(cls.evt_col_metadata_node + colId, nodeName, nodeValue))

    @classmethod
    def _writeColumnData(cls, colId, count, data):
        # hottest point; called rows*columns times
        l = _writeStdOut(_makeDataEvent(cls.evt_cell_data + colId, count))
        l += _writeStdOut(data)
        _writeStdOut(pack(_padding(l)))

    @classmethod
    def _writeDbMetadataAttr(cls, dbId, nodeName, attrName, attrValue):
        _writeStdOut(_make3StringEvent(cls.evt_db_metadata_node_attr + dbId, nodeName, attrName, attrValue))

    @classmethod
    def _writeTableMetadataAttr(cls, tblId, nodeName, attrName, attrValue):
        _writeStdOut(_make3StringEvent(cls.evt_tbl_metadata_node_attr + tblId, nodeName, attrName, attrValue))

    @classmethod
    def _writeColumnMetadataAttr(cls, colId, nodeName, attrName, attrValue):
        _writeStdOut(_make3StringEvent(cls.evt_col_metadata_node_attr + colId, nodeName, attrName, attrValue))

    @classmethod
    def _writeNextRow(cls, tableId):
        # called once per row
        _writeStdOut(_makeSimpleEvent(cls.evt_next_row + tableId))

    def writeDbMetadata(self, nodeName, nodeValue):
        """ this only supports writing to the default database """
        GeneralWriter._writeDbMetadata(0, nodeName.encode('ascii'), nodeValue.encode('utf-8'))

    def writeTableMetadata(self, table, nodeName, nodeValue):
        GeneralWriter._writeTableMetadata(table['_tableId'], nodeName.encode('ascii'), nodeValue.encode('utf-8'))

    def writeColumnMetadata(self, column, nodeName, nodeValue):
        GeneralWriter._writeColumnMetadata(column['_columnId'], nodeName.encode('ascii'), nodeValue.encode('utf-8'))

    def writeDbMetadataNodeAttr(self, nodeName, attrName, attrValue):
        """ this only supports writing to the default database """
        GeneralWriter._writeDbMetadataAttr(0, nodeName.encode('ascii'), attrName.encode('ascii'), attrValue.encode('utf-8'))

    def writeTableMetadataNodeAttr(self, table, nodeName, attrName, attrValue):
        GeneralWriter._writeTableMetadataAttr(table['_tableId'], nodeName.encode('ascii'), attrName.encode('ascii'), attrValue.encode('utf-8'))

    def writeColumnMetadataNodeAttr(self, column, nodeName, attrName, attrValue):
        GeneralWriter._writeColumnMetadataAttr(column['_columnId'], nodeName.encode('ascii'), attrName.encode('ascii'), attrValue.encode('utf-8'))

    def __init__(self, fileName, schemaFileName, schemaDbSpec, softwareName, versionString, tbl):
        """ Construct a General Writer object

            writer may be None if no actual output is desired
                else writer is expected to have a callable write attribute

            fileName is a string with name of the database that will be created

            schemaFileName is a string with the path to the file containing the schema

            schemaDbSpec is a string with the schema name of the database that will be created

            softwareName is a string

            versionString is a three-part number like "2.1.5"
        """

        GeneralWriter._writeHeader(fileName.encode('utf-8')
            , schemaFileName.encode('utf-8')
            , schemaDbSpec.encode('ascii'))

        GeneralWriter._writeSoftwareName(softwareName.encode('utf-8'), versionString.encode('ascii'))

        tableId = 0
        columnId = 0
        for t in tbl:
            tableId += 1
            GeneralWriter._writeNewTable(tableId, t.encode('ascii'))
            cols = tbl[t]
            for c in cols:
                columnId = columnId + 1
                cols[c]['_columnId'] = columnId
                expression = cols[c]['expression'] if 'expression' in cols[c] else c
                bits = cols[c]['elem_bits']
                GeneralWriter._writeNewColumn(columnId, tableId, bits, expression.encode('ascii'))
            tbl[t]['_tableId'] = tableId

        GeneralWriter._writeOpenStream()
        for t in tbl.values():
            for c in t.values():
                try:
                    if 'default' in c:
                        try:
                            GeneralWriter._writeColumnDefault(c['_columnId'], len(c['default']), c['default'])
                        except:
                            stderr.write("failed to set default for %s\n" % c)
                            raise
                except TypeError:
                    pass

    def __del__(self):
        try: GeneralWriter._writeEndStream()
        except: pass


if __name__ == "__main__":
    row = 0

    def performance():
        from array import array
        import random

        def gen_phred_range():
            for c in range(3, 41):
                yield chr(33+c)
        phred_range = list(gen_phred_range())
        reads = list([''.join(random.choices(('A', 'C', 'G', 'T'), k=300))] * 50000)
        quals = list([''.join(random.choices(phred_range, k=300))] * 50000)

        def getRead():
            return random.choice(reads).encode('ascii')

        def getQual():
            return random.choice(quals).encode('ascii')

        spec = {
            'SEQUENCE': {
                'READ': {
                    'expression': '(INSDC:dna:text)READ',
                    'elem_bits': 8,
                    'data': getRead
                },
                'QUALITY': {
                    'expression': '(INSDC:quality:phred)QUALITY',
                    'elem_bits': 8,
                    'data': getQual
                },
                'LABEL': {
                    'elem_bits': 8,
                    'default': 'templatecomplement'.encode('ascii')
                },
                'LABEL_START': {
                    'elem_bits': 32,
                },
                'LABEL_LENGTH': {
                    'elem_bits': 32,
                },
                'READ_START': {
                    'elem_bits': 32,
                },
                'READ_LENGTH': {
                    'elem_bits': 32,
                },
            },
        }
        gw = GeneralWriter('file name', 'test/bogus.vschema', 'Test:BOGUS:tbl', 'GeneralWriter-test', '1.0.0', spec)

        spec['SEQUENCE']['READ_LENGTH']['data'] = array('I', [150, 150])
        for _ in range(5_000_000):
            gw.write(spec['SEQUENCE'])

        gw = None


    def sanity():
        from array import array

        readData = [ "ACGT", "GTAACGT" ]

        def getRead():
            global row
            cur = row
            row = row + 1
            return readData[cur].encode('ascii')

        spec = {
            'SEQUENCE': {
                'READ': {
                    'expression': '(INSDC:dna:text)READ',
                    'elem_bits': 8,
                    'data': getRead
                },
                'QUALITY': {
                    'expression': '(INSDC:quality:phred)QUALITY',
                    'elem_bits': 8,
                },
                'LABEL': {
                    'elem_bits': 8,
                    'default': 'templatecomplement'.encode('ascii')
                },
                'LABEL_START': {
                    'elem_bits': 32,
                },
                'LABEL_LENGTH': {
                    'elem_bits': 32,
                },
                'READ_START': {
                    'elem_bits': 32,
                },
                'READ_LENGTH': {
                    'elem_bits': 32,
                },
            },
            'CONSENSUS': {
                'READ': {
                    'expression': '(INSDC:dna:text)READ',
                    'elem_bits': 8,
                },
                'QUALITY': {
                    'expression': '(INSDC:quality:phred)QUALITY',
                    'elem_bits': 8,
                },
                'LABEL': {
                    'elem_bits': 8,
                    'default': '2DFull'.encode('ascii')
                },
                'LABEL_START': {
                    'elem_bits': 32,
                    'default': array('I', [ 0 ])
                },
                'LABEL_LENGTH': {
                    'elem_bits': 32,
                    'default': array('I', [ 2 ])
                }
            },
        }
        gw = GeneralWriter('file name', 'test/bogus.vschema', 'Test:BOGUS:tbl', 'GeneralWriter-test', '1.0.0', spec)

        spec['SEQUENCE']['READ_LENGTH']['data'] = array('I', [ 1, 2 ])
        gw.write(spec['SEQUENCE'])
        gw.write(spec['SEQUENCE'])

        gw.logMessage('log message')
        gw.errorMessage('error message')

        gw = None

    performance()