#!/usr/bin/env python
# This file is distributed under the terms of the 2-clause BSD License.
# Copyright (c) 2017-2018, Almar Klein

"""
Python implementation of the Binary Structured Data Format (BSDF).

BSDF is a binary format for serializing structured (scientific) data.
See http://bsdf.io for more information.

This is the reference implementation, which is relatively relatively
sophisticated, providing e.g. lazy loading of blobs and streamed
reading/writing. A simpler Python implementation is available as
``bsdf_lite.py``.

This module has no dependencies and works on Python 2.7 and 3.4+.

Note: on Legacy Python (Python 2.7), non-Unicode strings are encoded as bytes.
"""

# todo: in 2020, remove six stuff, __future__ and _isidentifier
# todo: in 2020, remove 'utf-8' args to encode/decode; it's faster

from __future__ import absolute_import, division, print_function

import bz2
import hashlib
import logging
import os
import re
import struct
import sys
import types
import zlib
from io import BytesIO

logger = logging.getLogger(__name__)

# Notes on versioning: the major and minor numbers correspond to the
# BSDF format version. The major number if increased when backward
# incompatible changes are introduced. An implementation must raise an
# exception when the file being read has a higher major version. The
# minor number is increased when new backward compatible features are
# introduced. An implementation must display a warning when the file
# being read has a higher minor version. The patch version is increased
# for subsequent releases of the implementation.
VERSION = 2, 1, 2
__version__ = '.'.join(str(i) for i in VERSION)


# %% The encoder and decoder implementation

# From six.py
PY3 = sys.version_info[0] >= 3
if PY3:
    text_type = str
    string_types = str
    unicode_types = str
    integer_types = int
    classtypes = type
else:  # pragma: no cover
    logging.basicConfig()  # avoid "no handlers found" error
    text_type = unicode  # noqa
    string_types = basestring  # noqa
    unicode_types = unicode  # noqa
    integer_types = (int, long)  # noqa
    classtypes = type, types.ClassType

# Shorthands
spack = struct.pack
strunpack = struct.unpack


def lencode(x):
    """ Encode an unsigned integer into a variable sized blob of bytes.
    """
    # We could support 16 bit and 32 bit as well, but the gain is low, since
    # 9 bytes for collections with over 250 elements is marginal anyway.
    if x <= 250:
        return spack('<B', x)
    # elif x < 65536:
    #     return spack('<BH', 251, x)
    # elif x < 4294967296:
    #     return spack('<BI', 252, x)
    else:
        return spack('<BQ', 253, x)


# Include len decoder for completeness; we've inlined it for performance.
def lendecode(f):
    """ Decode an unsigned integer from a file.
    """
    n = strunpack('<B', f.read(1))[0]
    if n == 253: n = strunpack('<Q', f.read(8))[0]  # noqa
    return n


def encode_type_id(b, ext_id):
    """ Encode the type identifier, with or without extension id.
    """
    if ext_id is not None:
        bb = ext_id.encode('UTF-8')
        return b.upper() + lencode(len(bb)) + bb  # noqa
    else:
        return b  # noqa


def _isidentifier(s):  # pragma: no cover
    """ Use of str.isidentifier() for Legacy Python, but slower.
    """
    # http://stackoverflow.com/questions/2544972/
    return (isinstance(s, string_types) and
            re.match(r'^\w+$', s, re.UNICODE) and
            re.match(r'^[0-9]', s) is None)


class BsdfSerializer(object):
    """ Instances of this class represent a BSDF encoder/decoder.

    It acts as a placeholder for a set of extensions and encoding/decoding
    options. Use this to predefine extensions and options for high
    performance encoding/decoding. For general use, see the functions
    `save()`, `encode()`, `load()`, and `decode()`.

    This implementation of BSDF supports streaming lists (keep adding
    to a list after writing the main file), lazy loading of blobs, and
    in-place editing of blobs (for streams opened with a+).

    Options for encoding:

    * compression (int or str): ``0`` or "no" for no compression (default),
      ``1`` or "zlib" for Zlib compression (same as zip files and PNG), and
      ``2`` or "bz2" for Bz2 compression (more compact but slower writing).
      Note that some BSDF implementations (e.g. JavaScript) may not support
      compression.
    * use_checksum (bool): whether to include a checksum with binary blobs.
    * float64 (bool): Whether to write floats as 64 bit (default) or 32 bit.

    Options for decoding:

    * load_streaming (bool): if True, and the final object in the structure was
      a stream, will make it available as a stream in the decoded object.
    * lazy_blob (bool): if True, bytes are represented as Blob objects that can
      be used to lazily access the data, and also overwrite the data if the
      file is open in a+ mode.
    """

    def __init__(self, extensions=None, **options):
        self._extensions = {}  # name -> extension
        self._extensions_by_cls = {}  # cls -> (name, extension.encode)
        if extensions is None:
            extensions = standard_extensions
        for extension in extensions:
            self.add_extension(extension)
        self._parse_options(**options)

    def _parse_options(self,
                       compression=0, use_checksum=False, float64=True,
                       load_streaming=False, lazy_blob=False):

        # Validate compression
        if isinstance(compression, string_types):
            m = {'no': 0, 'zlib': 1, 'bz2': 2}
            compression = m.get(compression.lower(), compression)
        if compression not in (0, 1, 2):
            raise TypeError('Compression must be 0, 1, 2, '
                            '"no", "zlib", or "bz2"')
        self._compression = compression

        # Other encoding args
        self._use_checksum = bool(use_checksum)
        self._float64 = bool(float64)

        # Decoding args
        self._load_streaming = bool(load_streaming)
        self._lazy_blob = bool(lazy_blob)

    def add_extension(self, extension_class):
        """ Add an extension to this serializer instance, which must be
        a subclass of Extension. Can be used as a decorator.
        """
        # Check class
        if not (isinstance(extension_class, type) and
                issubclass(extension_class, Extension)):
            raise TypeError('add_extension() expects a Extension class.')
        extension = extension_class()

        # Get name
        name = extension.name
        if not isinstance(name, str):
            raise TypeError('Extension name must be str.')
        if len(name) == 0 or len(name) > 250:
            raise NameError('Extension names must be nonempty and shorter '
                            'than 251 chars.')
        if name in self._extensions:
            logger.warn('BSDF warning: overwriting extension "%s", '
                        'consider removing first' % name)

        # Get classes
        cls = extension.cls
        if not cls:
            clss = []
        elif isinstance(cls, (tuple, list)):
            clss = cls
        else:
            clss = [cls]
        for cls in clss:
            if not isinstance(cls, classtypes):
                raise TypeError('Extension classes must be types.')

        # Store
        for cls in clss:
            self._extensions_by_cls[cls] = name, extension.encode
        self._extensions[name] = extension
        return extension_class

    def remove_extension(self, name):
        """ Remove a converted by its unique name.
        """
        if not isinstance(name, str):
            raise TypeError('Extension name must be str.')
        if name in self._extensions:
            self._extensions.pop(name)
        for cls in list(self._extensions_by_cls.keys()):
            if self._extensions_by_cls[cls][0] == name:
                self._extensions_by_cls.pop(cls)

    def _encode(self, f, value, streams, ext_id):
        """ Main encoder function.
        """
        x = encode_type_id

        if value is None:
            f.write(x(b'v', ext_id))  # V for void
        elif value is True:
            f.write(x(b'y', ext_id))  # Y for yes
        elif value is False:
            f.write(x(b'n', ext_id))  # N for no
        elif isinstance(value, integer_types):
            if -32768 <= value <= 32767:
                f.write(x(b'h', ext_id) + spack('h', value))  # H for ...
            else:
                f.write(x(b'i', ext_id) + spack('<q', value))  # I for int
        elif isinstance(value, float):
            if self._float64:
                f.write(x(b'd', ext_id) + spack('<d', value))  # D for double
            else:
                f.write(x(b'f', ext_id) + spack('<f', value))  # f for float
        elif isinstance(value, unicode_types):
            bb = value.encode('UTF-8')
            f.write(x(b's', ext_id) + lencode(len(bb)))  # S for str
            f.write(bb)
        elif isinstance(value, (list, tuple)):
            f.write(x(b'l', ext_id) + lencode(len(value)))  # L for list
            for v in value:
                self._encode(f, v, streams, None)
        elif isinstance(value, dict):
            f.write(x(b'm', ext_id) + lencode(len(value)))  # M for mapping
            for key, v in value.items():
                if PY3:
                    assert key.isidentifier()  # faster
                else:  # pragma: no cover
                    assert _isidentifier(key)
                # yield ' ' * indent + key
                name_b = key.encode('UTF-8')
                f.write(lencode(len(name_b)))
                f.write(name_b)
                self._encode(f, v, streams, None)
        elif isinstance(value, bytes):
            f.write(x(b'b', ext_id))  # B for blob
            blob = Blob(value, compression=self._compression,
                        use_checksum=self._use_checksum)
            blob._to_file(f)  # noqa
        elif isinstance(value, Blob):
            f.write(x(b'b', ext_id))  # B for blob
            value._to_file(f)  # noqa
        elif isinstance(value, BaseStream):
            # Initialize the stream
            if value.mode != 'w':
                raise ValueError('Cannot serialize a read-mode stream.')
            elif isinstance(value, ListStream):
                f.write(x(b'l', ext_id) + spack('<BQ', 255, 0))  # L for list
            else:
                raise TypeError('Only ListStream is supported')
            # Mark this as *the* stream, and activate the stream.
            # The save() function verifies this is the last written object.
            if len(streams) > 0:
                raise ValueError('Can only have one stream per file.')
            streams.append(value)
            value._activate(f, self._encode, self._decode)  # noqa
        else:
            if ext_id is not None:
                raise ValueError(
                    'Extension %s wronfully encodes object to another '
                    'extension object (though it may encode to a list/dict '
                    'that contains other extension objects).' % ext_id)
            # Try if the value is of a type we know
            ex = self._extensions_by_cls.get(value.__class__, None)
            # Maybe its a subclass of a type we know
            if ex is None:
                for name, c in self._extensions.items():
                    if c.match(self, value):
                        ex = name, c.encode
                        break
                else:
                    ex = None
            # Success or fail
            if ex is not None:
                ext_id2, extension_encode = ex
                self._encode(f, extension_encode(self, value),
                             streams, ext_id2)
            else:
                t = ('Class %r is not a valid base BSDF type, nor is it '
                     'handled by an extension.')
                raise TypeError(t % value.__class__.__name__)

    def _decode(self, f):
        """ Main decoder function.
        """

        # Get value
        char = f.read(1)
        c = char.lower()

        # Conversion (uppercase value identifiers signify converted values)
        if not char:
            raise EOFError()
        elif char != c:
            n = strunpack('<B', f.read(1))[0]
            # if n == 253: n = strunpack('<Q', f.read(8))[0]  # noqa - noneed
            ext_id = f.read(n).decode('UTF-8')
        else:
            ext_id = None

        if c == b'v':
            value = None
        elif c == b'y':
            value = True
        elif c == b'n':
            value = False
        elif c == b'h':
            value = strunpack('<h', f.read(2))[0]
        elif c == b'i':
            value = strunpack('<q', f.read(8))[0]
        elif c == b'f':
            value = strunpack('<f', f.read(4))[0]
        elif c == b'd':
            value = strunpack('<d', f.read(8))[0]
        elif c == b's':
            n_s = strunpack('<B', f.read(1))[0]
            if n_s == 253: n_s = strunpack('<Q', f.read(8))[0]  # noqa
            value = f.read(n_s).decode('UTF-8')
        elif c == b'l':
            n = strunpack('<B', f.read(1))[0]
            if n >= 254:
                # Streaming
                closed = n == 254
                n = strunpack('<Q', f.read(8))[0]
                if self._load_streaming:
                    value = ListStream(n if closed else 'r')
                    value._activate(f, self._encode, self._decode)  # noqa
                elif closed:
                    value = [self._decode(f) for i in range(n)]
                else:
                    value = []
                    try:
                        while True:
                            value.append(self._decode(f))
                    except EOFError:
                        pass
            else:
                # Normal
                if n == 253: n = strunpack('<Q', f.read(8))[0]  # noqa
                value = [self._decode(f) for i in range(n)]
        elif c == b'm':
            value = dict()
            n = strunpack('<B', f.read(1))[0]
            if n == 253: n = strunpack('<Q', f.read(8))[0]  # noqa
            for i in range(n):
                n_name = strunpack('<B', f.read(1))[0]
                if n_name == 253: n_name = strunpack('<Q', f.read(8))[0]  # noqa
                assert n_name > 0
                name = f.read(n_name).decode('UTF-8')
                value[name] = self._decode(f)
        elif c == b'b':
            if self._lazy_blob:
                value = Blob((f, True))
            else:
                blob = Blob((f, False))
                value = blob.get_bytes()
        else:
            raise RuntimeError('Parse error %r' % char)

        # Convert value if we have an extension for it
        if ext_id is not None:
            extension = self._extensions.get(ext_id, None)
            if extension is not None:
                value = extension.decode(self, value)
            else:
                logger.warn('BSDF warning: no extension found for %r' % ext_id)

        return value

    def encode(self, ob):
        """ Save the given object to bytes.
        """
        f = BytesIO()
        self.save(f, ob)
        return f.getvalue()

    def save(self, f, ob):
        """ Write the given object to the given file object.
        """
        f.write(b'BSDF')
        f.write(struct.pack('<B', VERSION[0]))
        f.write(struct.pack('<B', VERSION[1]))

        # Prepare streaming, this list will have 0 or 1 item at the end
        streams = []

        self._encode(f, ob, streams, None)

        # Verify that stream object was at the end, and add initial elements
        if len(streams) > 0:
            stream = streams[0]
            if stream._start_pos != f.tell():
                raise ValueError('The stream object must be '
                                 'the last object to be encoded.')

    def decode(self, bb):
        """ Load the data structure that is BSDF-encoded in the given bytes.
        """
        f = BytesIO(bb)
        return self.load(f)

    def load(self, f):
        """ Load a BSDF-encoded object from the given file object.
        """
        # Check magic string
        f4 = f.read(4)
        if f4 != b'BSDF':
            raise RuntimeError('This does not look like a BSDF file: %r' % f4)
        # Check version
        major_version = strunpack('<B', f.read(1))[0]
        minor_version = strunpack('<B', f.read(1))[0]
        file_version = '%i.%i' % (major_version, minor_version)
        if major_version != VERSION[0]:  # major version should be 2
            t = ('Reading file with different major version (%s) '
                 'from the implementation (%s).')
            raise RuntimeError(t % (__version__, file_version))
        if minor_version > VERSION[1]:  # minor should be < ours
            t = ('BSDF warning: reading file with higher minor version (%s) '
                 'than the implementation (%s).')
            logger.warn(t % (__version__, file_version))

        return self._decode(f)


# %% Streaming and blob-files


class BaseStream(object):
    """ Base class for streams.
    """

    def __init__(self, mode='w'):
        self._i = 0
        self._count = -1
        if isinstance(mode, int):
            self._count = mode
            mode = 'r'
        elif mode == 'w':
            self._count = 0
        assert mode in ('r', 'w')
        self._mode = mode
        self._f = None
        self._start_pos = 0

    def _activate(self, file, encode_func, decode_func):
        if self._f is not None:  # Associated with another write
            raise IOError('Stream object cannot be activated twice?')
        self._f = file
        self._start_pos = self._f.tell()
        self._encode = encode_func
        self._decode = decode_func

    @property
    def mode(self):
        """ The mode of this stream: 'r' or 'w'.
        """
        return self._mode


class ListStream(BaseStream):
    """ A streamable list object used for writing or reading.
    In read mode, it can also be iterated over.
    """

    @property
    def count(self):
        """ The number of elements in the stream (can be -1 for unclosed
        streams in read-mode).
        """
        return self._count

    @property
    def index(self):
        """ The current index of the element to read/write.
        """
        return self._i

    def append(self, item):
        """ Append an item to the streaming list. The object is immediately
        serialized and written to the underlying file.
        """
        # if self._mode != 'w':
        #     raise IOError('This ListStream is not in write mode.')
        if self._count != self._i:
            raise IOError('Can only append items to the end of the stream.')
        if self._f is None:
            raise IOError('List stream is not associated with a file yet.')
        if self._f.closed:
            raise IOError('Cannot stream to a close file.')
        self._encode(self._f, item, [self], None)
        self._i += 1
        self._count += 1

    def close(self, unstream=False):
        """ Close the stream, marking the number of written elements. New
        elements may still be appended, but they won't be read during decoding.
        If ``unstream`` is False, the stream is turned into a regular list
        (not streaming).
        """
        # if self._mode != 'w':
        #     raise IOError('This ListStream is not in write mode.')
        if self._count != self._i:
            raise IOError('Can only close when at the end of the stream.')
        if self._f is None:
            raise IOError('ListStream is not associated with a file yet.')
        if self._f.closed:
            raise IOError('Cannot close a stream on a close file.')
        i = self._f.tell()
        self._f.seek(self._start_pos - 8 - 1)
        self._f.write(spack('<B', 253 if unstream else 254))
        self._f.write(spack('<Q', self._count))
        self._f.seek(i)

    def next(self):
        """ Read and return the next element in the streaming list.
        Raises StopIteration if the stream is exhausted.
        """
        if self._mode != 'r':
            raise IOError('This ListStream in not in read mode.')
        if self._f is None:
            raise IOError('ListStream is not associated with a file yet.')
        if getattr(self._f, 'closed', None):  # not present on 2.7 http req :/
            raise IOError('Cannot read a stream from a close file.')
        if self._count >= 0:
            if self._i >= self._count:
                raise StopIteration()
            self._i += 1
            return self._decode(self._f)
        else:
            # This raises EOFError at some point.
            try:
                res = self._decode(self._f)
                self._i += 1
                return res
            except EOFError:
                self._count = self._i
                raise StopIteration()

    def __iter__(self):
        if self._mode != 'r':
            raise IOError('Cannot iterate: ListStream in not in read mode.')
        return self

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


class Blob(object):
    """ Object to represent a blob of bytes. When used to write a BSDF file,
    it's a wrapper for bytes plus properties such as what compression to apply.
    When used to read a BSDF file, it can be used to read the data lazily, and
    also modify the data if reading in 'r+' mode and the blob isn't compressed.
    """

    # For now, this does not allow re-sizing blobs (within the allocated size)
    # but this can be added later.

    def __init__(self, bb, compression=0, extra_size=0, use_checksum=False):
        if isinstance(bb, bytes):
            self._f = None
            self.compressed = self._from_bytes(bb, compression)
            self.compression = compression
            self.allocated_size = self.used_size + extra_size
            self.use_checksum = use_checksum
        elif isinstance(bb, tuple) and len(bb) == 2 and hasattr(bb[0], 'read'):
            self._f, allow_seek = bb
            self.compressed = None
            self._from_file(self._f, allow_seek)
            self._modified = False
        else:
            raise TypeError('Wrong argument to create Blob.')

    def _from_bytes(self, value, compression):
        """ When used to wrap bytes in a blob.
        """
        if compression == 0:
            compressed = value
        elif compression == 1:
            compressed = zlib.compress(value, 9)
        elif compression == 2:
            compressed = bz2.compress(value, 9)
        else:  # pragma: no cover
            assert False, 'Unknown compression identifier'

        self.data_size = len(value)
        self.used_size = len(compressed)
        return compressed

    def _to_file(self, f):
        """ Private friend method called by encoder to write a blob to a file.
        """
        # Write sizes - write at least in a size that allows resizing
        if self.allocated_size <= 250 and self.compression == 0:
            f.write(spack('<B', self.allocated_size))
            f.write(spack('<B', self.used_size))
            f.write(lencode(self.data_size))
        else:
            f.write(spack('<BQ', 253, self.allocated_size))
            f.write(spack('<BQ', 253, self.used_size))
            f.write(spack('<BQ', 253, self.data_size))
        # Compression and checksum
        f.write(spack('B', self.compression))
        if self.use_checksum:
            f.write(b'\xff' + hashlib.md5(self.compressed).digest())
        else:
            f.write(b'\x00')
        # Byte alignment (only necessary for uncompressed data)
        if self.compression == 0:
            alignment = 8 - (f.tell() + 1) % 8  # +1 for the byte to write
            f.write(spack('<B', alignment))  # padding for byte alignment
            f.write(b'\x00' * alignment)
        else:
            f.write(spack('<B', 0))
        # The actual data and extra space
        f.write(self.compressed)
        f.write(b'\x00' * (self.allocated_size - self.used_size))

    def _from_file(self, f, allow_seek):
        """ Used when a blob is read by the decoder.
        """
        # Read blob header data (5 to 42 bytes)
        # Size
        allocated_size = strunpack('<B', f.read(1))[0]
        if allocated_size == 253: allocated_size = strunpack('<Q', f.read(8))[0]  # noqa
        used_size = strunpack('<B', f.read(1))[0]
        if used_size == 253: used_size = strunpack('<Q', f.read(8))[0]  # noqa
        data_size = strunpack('<B', f.read(1))[0]
        if data_size == 253: data_size = strunpack('<Q', f.read(8))[0]  # noqa
        # Compression and checksum
        compression = strunpack('<B', f.read(1))[0]
        has_checksum = strunpack('<B', f.read(1))[0]
        if has_checksum:
            checksum = f.read(16)
        # Skip alignment
        alignment = strunpack('<B', f.read(1))[0]
        f.read(alignment)
        # Get or skip data + extra space
        if allow_seek:
            self.start_pos = f.tell()
            self.end_pos = self.start_pos + used_size
            f.seek(self.start_pos + allocated_size)
        else:
            self.start_pos = None
            self.end_pos = None
            self.compressed = f.read(used_size)
            f.read(allocated_size - used_size)
        # Store info
        self.alignment = alignment
        self.compression = compression
        self.use_checksum = checksum if has_checksum else None
        self.used_size = used_size
        self.allocated_size = allocated_size
        self.data_size = data_size

    def seek(self, p):
        """ Seek to the given position (relative to the blob start).
        """
        if self._f is None:
            raise RuntimeError('Cannot seek in a blob '
                               'that is not created by the BSDF decoder.')
        if p < 0:
            p = self.allocated_size + p
        if p < 0 or p > self.allocated_size:
            raise IOError('Seek beyond blob boundaries.')
        self._f.seek(self.start_pos + p)

    def tell(self):
        """ Get the current file pointer position (relative to the blob start).
        """
        if self._f is None:
            raise RuntimeError('Cannot tell in a blob '
                               'that is not created by the BSDF decoder.')
        return self._f.tell() - self.start_pos

    def write(self, bb):
        """ Write bytes to the blob.
        """
        if self._f is None:
            raise RuntimeError('Cannot write in a blob '
                               'that is not created by the BSDF decoder.')
        if self.compression:
            raise IOError('Cannot arbitrarily write in compressed blob.')
        if self._f.tell() + len(bb) > self.end_pos:
            raise IOError('Write beyond blob boundaries.')
        self._modified = True
        return self._f.write(bb)

    def read(self, n):
        """ Read n bytes from the blob.
        """
        if self._f is None:
            raise RuntimeError('Cannot read in a blob '
                               'that is not created by the BSDF decoder.')
        if self.compression:
            raise IOError('Cannot arbitrarily read in compressed blob.')
        if self._f.tell() + n > self.end_pos:
            raise IOError('Read beyond blob boundaries.')
        return self._f.read(n)

    def get_bytes(self):
        """ Get the contents of the blob as bytes.
        """
        if self.compressed is not None:
            compressed = self.compressed
        else:
            i = self._f.tell()
            self.seek(0)
            compressed = self._f.read(self.used_size)
            self._f.seek(i)
        if self.compression == 0:
            value = compressed
        elif self.compression == 1:
            value = zlib.decompress(compressed)
        elif self.compression == 2:
            value = bz2.decompress(compressed)
        else:  # pragma: no cover
            raise RuntimeError('Invalid compression %i' % self.compression)
        return value

    def update_checksum(self):
        """ Reset the blob's checksum if present. Call this after modifying
        the data.
        """
        # or ... should the presence of a checksum mean that data is proteced?
        if self.use_checksum and self._modified:
            self.seek(0)
            compressed = self._f.read(self.used_size)
            self._f.seek(self.start_pos - self.alignment - 1 - 16)
            self._f.write(hashlib.md5(compressed).digest())


# %% High-level functions


def encode(ob, extensions=None, **options):
    """ Save (BSDF-encode) the given object to bytes.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    return s.encode(ob)


def save(f, ob, extensions=None, **options):
    """ Save (BSDF-encode) the given object to the given filename or
    file object. See` BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    if isinstance(f, string_types):
        with open(f, 'wb') as fp:
            return s.save(fp, ob)
    else:
        return s.save(f, ob)


def decode(bb, extensions=None, **options):
    """ Load a (BSDF-encoded) structure from bytes.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    return s.decode(bb)


def load(f, extensions=None, **options):
    """ Load a (BSDF-encoded) structure from the given filename or file object.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    if isinstance(f, string_types):
        if f.startswith(('~/', '~\\')):  # pragma: no cover
            f = os.path.expanduser(f)
        with open(f, 'rb') as fp:
            return s.load(fp)
    else:
        return s.load(f)


# Aliases for json compat
loads = decode
dumps = encode


# %% Standard extensions

# Defining extensions as a dict would be more compact and feel lighter, but
# that would only allow lambdas, which is too limiting, e.g. for ndarray
# extension.

class Extension(object):
    """ Base class to implement BSDF extensions for special data types.

    Extension classes are provided to the BSDF serializer, which
    instantiates the class. That way, the extension can be somewhat dynamic:
    e.g. the NDArrayExtension exposes the ndarray class only when numpy
    is imported.

    A extension instance must have two attributes. These can be attribiutes of
    the class, or of the instance set in ``__init__()``:

    * name (str): the name by which encoded values will be identified.
    * cls (type): the type (or list of types) to match values with.
      This is optional, but it makes the encoder select extensions faster.

    Further, it needs 3 methods:

    * `match(serializer, value) -> bool`: return whether the extension can
      convert the given value. The default is ``isinstance(value, self.cls)``.
    * `encode(serializer, value) -> encoded_value`: the function to encode a
      value to more basic data types.
    * `decode(serializer, encoded_value) -> value`: the function to decode an
      encoded value back to its intended representation.

    """

    name = ''
    cls = ()

    def __repr__(self):
        return '<BSDF extension %r at 0x%s>' % (self.name, hex(id(self)))

    def match(self, s, v):
        return isinstance(v, self.cls)

    def encode(self, s, v):
        raise NotImplementedError()

    def decode(self, s, v):
        raise NotImplementedError()


class ComplexExtension(Extension):

    name = 'c'
    cls = complex

    def encode(self, s, v):
        return (v.real, v.imag)

    def decode(self, s, v):
        return complex(v[0], v[1])


class NDArrayExtension(Extension):

    name = 'ndarray'

    def __init__(self):
        if 'numpy' in sys.modules:
            import numpy as np
            self.cls = np.ndarray

    def match(self, s, v):  # pragma: no cover - e.g. work for nd arrays in JS
        return (hasattr(v, 'shape') and
                hasattr(v, 'dtype') and
                hasattr(v, 'tobytes'))

    def encode(self, s, v):
        return dict(shape=v.shape,
                    dtype=text_type(v.dtype),
                    data=v.tobytes())

    def decode(self, s, v):
        try:
            import numpy as np
        except ImportError:  # pragma: no cover
            return v
        a = np.frombuffer(v['data'], dtype=v['dtype'])
        a.shape = v['shape']
        return a


standard_extensions = [ComplexExtension, NDArrayExtension]


if __name__ == '__main__':
    # Invoke CLI
    import bsdf_cli
    bsdf_cli.main()
