# Copyright (c) 2014-2016 Siphon Contributors.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
"""Handle binary stream returns in NCStream format."""

from collections import OrderedDict
import itertools
import logging
import zlib

import numpy as np

from . import cdmrfeature_pb2 as cdmrf
from . import ncStream_pb2 as stream  # noqa

MAGIC_HEADER = b'\xad\xec\xce\xda'
MAGIC_DATA = b'\xab\xec\xce\xba'
MAGIC_DATA2 = b'\xab\xeb\xbe\xba'
MAGIC_VDATA = b'\xab\xef\xfe\xba'
MAGIC_VEND = b'\xed\xef\xfe\xda'
MAGIC_ERR = b'\xab\xad\xba\xda'

MAGIC_HEADERCOV = b'\xad\xed\xde\xda'
MAGIC_DATACOV = b'\xab\xed\xde\xba'

log = logging.getLogger(__name__)
log.setLevel(logging.WARNING)


#
# NCStream handling
#
def read_ncstream_data(fobj):
    """Handle reading an NcStream v1 data block from a file-like object."""
    data = read_proto_object(fobj, stream.Data)
    if data.dataType in (stream.STRING, stream.OPAQUE) or data.vdata:
        log.debug('Reading string/opaque/vlen')
        num_obj = read_var_int(fobj)
        log.debug('Num objects: %d', num_obj)
        blocks = [read_block(fobj) for _ in range(num_obj)]
        if data.dataType == stream.STRING:
            blocks = [b.decode('utf-8', errors='ignore') for b in blocks]

        # Again endian isn't coded properly
        dt = data_type_to_numpy(data.dataType).newbyteorder('>')
        if data.vdata:
            return np.array([np.frombuffer(b, dtype=dt) for b in blocks], dtype=object)
        else:
            return np.array(blocks, dtype=dt)
    elif data.dataType in _dtype_lookup:
        log.debug('Reading array data')
        bin_data = read_block(fobj)
        log.debug('Binary data: %s', bin_data)

        # Hard code to big endian for now since it's not encoded correctly
        dt = data_type_to_numpy(data.dataType).newbyteorder('>')

        # Handle decompressing the bytes
        if data.compress == stream.DEFLATE:
            bin_data = zlib.decompress(bin_data)
            assert len(bin_data) == data.uncompressedSize
        elif data.compress != stream.NONE:
            raise NotImplementedError(f'Compression type {data.compress} not implemented!')

        # Turn bytes into an array
        return reshape_array(data, np.frombuffer(bin_data, dtype=dt))
    elif data.dataType == stream.STRUCTURE:
        sd = read_proto_object(fobj, stream.StructureData)

        # Make a datatype appropriate to the rows of struct
        endian = '>' if data.bigend else '<'
        dt = np.dtype([(endian, np.void, sd.rowLength)])

        # Turn bytes into an array
        return reshape_array(data, np.frombuffer(sd.data, dtype=dt))
    elif data.dataType == stream.SEQUENCE:
        log.debug('Reading sequence')
        blocks = []
        magic = read_magic(fobj)
        while magic != MAGIC_VEND:
            if magic == MAGIC_VDATA:
                log.error('Bad magic for struct/seq data!')
            blocks.append(read_proto_object(fobj, stream.StructureData))
            magic = read_magic(fobj)
        return data, blocks
    else:
        raise NotImplementedError(f"Don't know how to handle data type: {data.dataType}")


def read_ncstream_data2(fobj):
    """Handle reading an NcStream v2 data block from a file-like object."""
    data = read_proto_object(fobj, stream.DataCol)
    return datacol_to_array(data)


def read_ncstream_err(fobj):
    """Handle reading an NcStream error from a file-like object and raise as error."""
    err = read_proto_object(fobj, stream.Error)
    raise RuntimeError(err.message)


ncstream_table = {MAGIC_HEADER: lambda f: read_proto_object(f, stream.Header),
                  MAGIC_DATA: read_ncstream_data,
                  MAGIC_DATA2: read_ncstream_data2,
                  MAGIC_ERR: read_ncstream_err}


def read_ncstream_messages(fobj):
    """Read a collection of NcStream messages from a file-like object."""
    return read_messages(fobj, ncstream_table)


#
# CDMRemoteFeature handling
#
cdmrf_table = {MAGIC_HEADERCOV: lambda f: read_proto_object(f, cdmrf.CoverageDataset),
               MAGIC_DATACOV: lambda f: read_proto_object(f, cdmrf.CoverageDataResponse),
               MAGIC_DATA2: read_ncstream_data2,  # For coordinates
               MAGIC_ERR: read_ncstream_err}


def read_cdmrf_messages(fobj):
    """Read a collection of CDMRemoteFeature messages from a file-like object."""
    return read_messages(fobj, cdmrf_table)


#
# General Utilities
#
def read_messages(fobj, magic_table):
    """Read messages from a file-like object until stream is exhausted."""
    messages = []

    while True:
        magic = read_magic(fobj)
        if not magic:
            break

        func = magic_table.get(magic)
        if func is not None:
            messages.append(func(fobj))
        else:
            log.error('Unknown magic: ' + str(' '.join(f'{b:02x}'
                                                       for b in bytearray(magic))))

    return messages


def read_proto_object(fobj, klass):
    """Read a block of data and parse using the given protobuf object."""
    log.debug('%s chunk', klass.__name__)
    obj = klass()
    obj.ParseFromString(read_block(fobj))
    log.debug('Header: %s', str(obj))
    return obj


def read_magic(fobj):
    """Read magic bytes.

    Parameters
    ----------
    fobj : file-like object
        The file to read from.

    Returns
    -------
    bytes
        magic byte sequence read

    """
    return fobj.read(4)


def read_block(fobj):
    """Read a block.

    Reads a block from a file object by first reading the number of bytes to read, which must
    be encoded as a variable-byte length integer.

    Parameters
    ----------
    fobj : file-like object
        The file to read from.

    Returns
    -------
    bytes
        block of bytes read

    """