#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import base64
import logging
import struct
import io

from gremlin_python.process.traversal import Traverser

try:
    import ujson as json
    if int(json.__version__[0]) < 2:
        logging.warning("Detected ujson version below 2.0.0. This is not recommended as precision may be lost.")
except ImportError:
    import json

from gremlin_python.structure.io import graphbinaryV4, graphsonV4
from gremlin_python.structure.io.util import Marker

__author__ = 'David M. Brown (davebshow@gmail.com), Lyndon Bauto (lyndonb@bitquilltech.com)'

"""
GraphSONV4
"""

class GraphSONSerializersV4(object):
    """
    Message serializer for GraphSON. Allow users to pass custom reader,
    writer, and version kwargs for custom serialization. Otherwise,
    use current GraphSON version as default.
    """

    # KEEP TRACK OF CURRENT DEFAULTS
    DEFAULT_READER_CLASS = graphsonV4.GraphSONReader
    DEFAULT_WRITER_CLASS = graphsonV4.GraphSONWriter
    DEFAULT_VERSION = b"application/vnd.gremlin-v4.0+json"

    def __init__(self, reader=None, writer=None, version=None):
        if not version:
            version = self.DEFAULT_VERSION
        self._version = version
        if not reader:
            reader = self.DEFAULT_READER_CLASS()
        self._graphson_reader = reader
        if not writer:
            writer = self.DEFAULT_WRITER_CLASS()
        self._graphson_writer = writer

    @property
    def version(self):
        """Read only property"""
        return self._version

    def serialize_message(self, request_message):
        message = self.build_message(request_message.fields, request_message.gremlin)
        return message

    def build_message(self, fields, gremlin):
        message = {
            'gremlin': gremlin  # gremlin wil always be a gremlin lang string script
        }
        for k, v in fields.items():
            message[k] = self._graphson_writer.to_dict(v)
        return self.finalize_message(message)

    def finalize_message(self, message):
        message = json.dumps(message)
        return message

    def deserialize_message(self, message, is_first_chunk=False):
        if is_first_chunk:
            msg = json.loads(message if isinstance(message, str) else message.decode('utf-8'))
            return self._graphson_reader.to_object(msg)
        else:
            # graphSON does not stream, all results are aggregated inside the first chunk
            return ""

"""
GraphBinaryV4
"""
class GraphBinarySerializersV4(object):
    DEFAULT_READER_CLASS = graphbinaryV4.GraphBinaryReader
    DEFAULT_WRITER_CLASS = graphbinaryV4.GraphBinaryWriter
    DEFAULT_VERSION = b"application/vnd.graphbinary-v4.0"

    max_int64 = 0xFFFFFFFFFFFFFFFF
    header_struct = struct.Struct('>b32sBQQ')
    header_pack = header_struct.pack
    int_pack = graphbinaryV4.int32_pack
    int32_unpack = struct.Struct(">i").unpack

    def __init__(self, reader=None, writer=None, version=None):
        if not version:
            version = self.DEFAULT_VERSION
        self._version = version
        if not reader:
            reader = self.DEFAULT_READER_CLASS()
        self._graphbinary_reader = reader
        if not writer:
            writer = self.DEFAULT_WRITER_CLASS()
        self._graphbinary_writer = writer
        self._bulked = False

    @property
    def version(self):
        """Read only property"""
        return self._version

    def serialize_message(self, request_message):
        message = self.build_message(request_message.fields, request_message.gremlin)
        return message

    def build_message(self, fields, gremlin):
        message = {
            'fields': fields,
            'gremlin': gremlin
        }
        return self.finalize_message(message)

    def finalize_message(self, message):
        ba = bytearray()

        ba.extend(graphbinaryV4.uint8_pack(0x84))
        fields = message["fields"]
        ba.extend(self.int_pack(len(fields)))
        for k, v in fields.items():
            self._graphbinary_writer.to_dict(k, ba)
            self._graphbinary_writer.to_dict(v, ba)

        gremlin = message['gremlin']
        # TODO: hack to remove type code from gremlin value for V4 message format, writer doesn't seem to have a way to
        #  write value directly by passing serializer types, check back when removing bytecode
        gremlin_ba = bytearray()
        self._graphbinary_writer.to_dict(gremlin, gremlin_ba)
        ba.extend(gremlin_ba[2:])

        return bytes(ba)

    def deserialize_message(self, message, is_first_chunk=False):
        if len(message) == 0:
            return {'status': {'code': 204},
                    'result': {'meta': {},
                               'data': []}}

        # for parsing string message via HTTP connections
        b = io.BytesIO(base64.b64decode(message) if isinstance(message, str) else message)

        if is_first_chunk:
            b.read(1)  # version
            self._bulked = b.read(1)[0] == 0x01

        result, readable = self.read_payload(b)
        if not readable:
            return {
                'result': {'meta': {},
                           'data': result}
            }
        status_code = self.int32_unpack(b.read(4))[0]  # status code
        status_msg = self._graphbinary_reader.to_object(b, graphbinaryV4.DataType.string, nullable=True)
        status_ex = self._graphbinary_reader.to_object(b, graphbinaryV4.DataType.string, nullable=True)
        # meta_attrs = self._graphbinary_reader.to_object(b, graphbinaryV4.DataType.map, nullable=False)

        b.close()

        msg = {'status': {'code': status_code,
                          'message': status_msg,
                          'exception': status_ex},
               'result': {'meta': {},
                          'data': result}}

        return msg

    def read_payload(self, buffer):
        results = []
        readable = True
        while buffer.readable():  # find method or way to access readable bytes without using buffer.getvalue()
            if buffer.tell() == len(buffer.getvalue()):
                readable = False
                break
            if self._bulked:
                item = self._graphbinary_reader.to_object(buffer)
                if item == Marker.end_of_stream():
                    self._bulked = False  # no more data expected, reset bulked flag
                    break
                bulk = self._graphbinary_reader.to_object(buffer)
                results.append(Traverser(item, bulk))
            else:
                data = self._graphbinary_reader.to_object(buffer)
                if data == Marker.end_of_stream():
                    break
                results.append(data)

        return results, readable