#!/usr/bin/env python3
#
# Serialization support for reading and writing RiskScape Types.
# This is basically the Python equivalent of the Type.toBytes() and
# Type.fromBytes() RiskScape Java code.

import os
import struct

# This makes it clear that the riskscape wrapper code generated the error
class RiskscapePyException(Exception):
    pass

# helper class for reading/writing numbers
class Number:
    @staticmethod
    # returns the appropriate format characters to use with struct.pack/.unpack APIs
    def get_pack_format(num_bytes, floating):
        # as per https://docs.python.org/3.4/library/struct.html#format-characters
        if floating:
            # only support 8-byte doubles
            pack_formats = {8: 'd', 4:'f'}
        else:
            pack_formats = {1: 'b', 2: 'h', 4: 'i', 8: 'q'}

        if num_bytes in pack_formats:
            # note that Java always uses big-endian ('>')
            return '>' + pack_formats[num_bytes]
        else:
            # developer error: we shouldn't hit this
            raise ValueError(num_bytes + "-byte numbers are not supported")

    @staticmethod
    def read(input_stream, num_bytes, floating=False):
        # read the data & unpack it, based on the type/size of number expected
        buffer = input_stream.read(num_bytes)
        pack_format = Number.get_pack_format(num_bytes, floating)
        # unpack always returns a tuple... *sigh*
        return struct.unpack(pack_format, buffer)[0]

    @staticmethod
    def write(output_stream, value, num_bytes, floating=False):
        pack_format = Number.get_pack_format(num_bytes, floating)
        output_stream.write(struct.pack(pack_format, value))


class Integer:
    @staticmethod
    def fromBytes(input_stream, unused):
        return Number.read(input_stream, num_bytes=8)

    @staticmethod
    def toBytes(output_stream, value, unused):
        # note the cast - python doesn't like it if we tell it to serialize
        # an int, but give it a float
        Number.write(output_stream, int(value), num_bytes=8)


class Enumeration:
    # enums are just encoded as 4-byte integers
    @staticmethod
    def fromBytes(input_stream, unused):
        return Number.read(input_stream, num_bytes=4)

    @staticmethod
    def toBytes(output_stream, value, unused):
        Number.write(output_stream, value, num_bytes=4)


class Floating:
    @staticmethod
    def fromBytes(input_stream, unused):
        return Number.read(input_stream, num_bytes=8, floating=True)

    @staticmethod
    def toBytes(output_stream, value, unused):
        Number.write(output_stream, float(value), num_bytes=8, floating=True)


class Smallfloat:
    @staticmethod
    def fromBytes(input_stream, unused):
        return Number.read(input_stream, num_bytes=4, floating=True)

    @staticmethod
    def toBytes(output_stream, value, unused):
        Number.write(output_stream, float(value), num_bytes=4, floating=True)

class Text:
    @staticmethod
    def fromBytes(input_stream, unused):
        # first read the 2-byte string-length
        length = Number.read(input_stream, num_bytes=2)

        # then read the UTF8-encoded text
        text = input_stream.read(length)
        return text.decode('utf8')

    @staticmethod
    def toBytes(output_stream, value, unused):
        # sanity-check in case the user didn't really gave us a string
        value = str(value)
        encoded_text = value.encode('utf8')
        Number.write(output_stream, len(encoded_text), num_bytes=2)
        # no need to use struct.pack here as encoding does the same job
        output_stream.write(encoded_text)


class Boolean:
    @staticmethod
    def fromBytes(input_stream, unused):
        return Number.read(input_stream, num_bytes=1) == 1

    @staticmethod
    def toBytes(output_stream, value, unused):
        if value not in [ True, False ]:
            raise TypeError(str(value) + " is not Boolean (True/False)")
        Number.write(output_stream, value == True, num_bytes=1)


class Nullable:
    @staticmethod
    def fromBytes(input_stream, contained_type):
        # first byte for Nullables is a boolean, and if true the contained type follows
        present = Boolean.fromBytes(input_stream, contained_type)
        return read(contained_type, input_stream) if present else None

    @staticmethod
    def toBytes(output_stream, value, contained_type):
        present = value != None
        Boolean.toBytes(output_stream, present, contained_type)
        if present:
            write(contained_type, output_stream, value)


class List:
    @staticmethod
    def fromBytes(input_stream, member_type):
        # first read the 4-byte list-length
        length = Number.read(input_stream, num_bytes=4)

        # then read the list members themselves
        values = []
        for i in range(length):
            values.append(read(member_type, input_stream))
        return values

    @staticmethod
    def toBytes(output_stream, values, member_type):
        Number.write(output_stream, len(values), num_bytes=4)
        for item in values:
            write(member_type, output_stream, item)


# NB - geometry is serialized as well known bytes - up to the user to manage the deserialization/marshalling
class Geometry:
    @staticmethod
    def fromBytes(input_stream, unused):
        # first read the 4-byte srid
        srid = Number.read(input_stream, num_bytes=4)
        # then read the 4-byte wkb 'packet' length - strictly speaking, wkb
        # doesn't require this, but it means we can read all the bytes without
        # understanding the protocol
        num_bytes = Number.read(input_stream, num_bytes=4)
        # User gets back a (blob 'o bytes, srid) tuple - the srid is mostly
        # meaningless, but can be used to at least send back the same to
        # mean the crs hasn't changed, which should avoid some unintentional
        # errors in case the returned geometry gets used for something  - future
        # changes might look to serialize the crs as part of the type OR
        # serialize the CRS as WKT (or just code?) when the the crs is missing
        # from the type
        return (input_stream.read(num_bytes), srid)

    @staticmethod
    def toBytes(output_stream, data, unused):

        if isinstance(data, tuple):
            try:
                byte_buffer = data[0]
                srid = data[1]
            except IndexError as err:
                # try to point the user in the right direction
                raise RiskscapePyException("Bad geometry tuple - %s.  Geometry "
                        "Geometry must be returned as a two member python tuple "
                        "e.g. (WKB, SRID)")
        else:
            # if it's not a tuple, we assume it's bytes
            srid = 0
            byte_buffer = bytes(data)

        Number.write(output_stream, srid, num_bytes=4)
        Number.write(output_stream, len(byte_buffer), num_bytes=4)
        output_stream.write(byte_buffer)


class Struct:
    @staticmethod
    def extract_type(remaining):
        nesting = 0
        index = 0
        for c in remaining:
            # we've reached a ',' boundary between struct members
            if c == ',' and nesting == 0:
                break
            elif c == '[':
                nesting += 1
            elif c == ']':
                nesting -= 1
            index += 1

        # split the string based on the member boundary found (if any)
        type_str = remaining[:index]
        remaining = remaining[index+1:]
        return (type_str.strip(), remaining.strip())

    @staticmethod
    def parse_members(members):
        parsed = []
        remaining = members
        while len(remaining) > 0:
            # first, extract the next 'member_name=>' from the remaining string
            member_name, remaining = remaining.split("=>", 1)
            # then split the remaining string up based on the type definition that follows
            type_str, remaining = Struct.extract_type(remaining)
            parsed.append((member_name, type_str))
        return parsed

    @staticmethod
    def fromBytes(input_stream, members):
        struct = {}
        for name, type_str in Struct.parse_members(members):
            struct[name] = read(type_str, input_stream)
        return struct

    @staticmethod
    def toBytes(output_stream, struct, members):
        for name, type_str in Struct.parse_members(members):
            try:
                write(type_str, output_stream, struct[name])
            except RiskscapePyException as e:
                raise RiskscapePyException("Problem found returning attribute '"
                     + name + "' in struct." + os.linesep + "\t" + str(e))
            except KeyError as ke:
                # the python KeyError is really not very informative
                raise RiskscapePyException("Problem found returning Struct - "
                    "attribute '" + name + "' was missing from the Python "
                    "dictionary returned by your function. The function "
                    "return-type in your INI file may not match your "
                    "python code.")


# Mappings of Types we have serialization support for
supported_types = {
    'Bool': Boolean,
    'Integer': Integer,
    'Enumeration': Enumeration,
    'Floating': Floating,
    'Smallfloat': Smallfloat,
    'List': List,
    'Nullable': Nullable,
    'Struct': Struct,
    'Text': Text,
    'Geometry': Geometry
    }

# returns the base type, along with with any contained type
# E.g. "Nullable[Text]" becomes: (Nullable, Text)
def split_contained_type(name):
    parts = name.split('[', 1)
    if len(parts) == 2:
        # wrapped type: Nullable, List, etc
        return (parts[0], parts[1].rstrip(']'))
    else:
        # simple type: Integer, Text, etc
        return (name, None)


# Looks up the Type-based serialization handler by name
def get_handler(type_name):
    if type_name not in supported_types:
        # ideally, we should catch this error in the RiskScape Java code
        raise RiskscapePyException('Serialization of ' + type_name +
                        ' is not supported for CPython. Try using Jython' +
                        ' or omitting this type from your function')
    return supported_types.get(type_name)


# Reads (i.e. deserializes) data from the input stream, based on the given type
def read(full_type, input_stream):
    (current_type, contained_type) = split_contained_type(full_type)
    return get_handler(current_type).fromBytes(input_stream, contained_type)


# Writes (i.e. serializes) the given type + value to the output stream
def write(full_type, output_stream, value):
    (current_type, contained_type) = split_contained_type(full_type)
    try:
        get_handler(current_type).toBytes(output_stream, value, contained_type)
        output_stream.flush()
    except (ValueError, TypeError) as e:
        tip = ""
        # give the user pointers for common mistakes
        if value is None:
            tip = ( os.linesep + "\tOne possible solution may be to define " +
                "this type as 'nullable' in your INI file return-type.")
        elif current_type == 'Struct':
            tip = ( os.linesep + "\tMake sure your python code returns " +
                "Structs in the format: { 'name1': value1, 'name2': value2 }")

        raise RiskscapePyException("Could not return value '" + str(value) +
            "' as a '" + current_type + "'. You may have specified the wrong " +
            "return-type in your function's INI file definition. " +
            "Details: " + str(e) + tip)
