# Copyright (c) Facebook, Inc. and its affiliates.
#
# Licensed 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.

# pyre-unsafe

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from nebula2.fbthrift.protocol.TProtocol import TProtocolBase, TProtocolException
from nebula2.fbthrift.Thrift import TType
import json, base64, sys

__all__ = ['TJSONProtocol', 'TJSONProtocolFactory']

VERSION = 1

COMMA = ','
COLON = ':'
LBRACE = '{'
RBRACE = '}'
LBRACKET = '['
RBRACKET = ']'
QUOTE = '"'
BACKSLASH = '\\'
ZERO = '0'

ESCSEQ = '\\u00'
ESCAPE_CHAR = '"\\bfnrt'
ESCAPE_CHAR_VALS = ['"', '\\', '\b', '\f', '\n', '\r', '\t']
NUMERIC_CHAR = '+-.0123456789Ee'

CTYPES = {TType.BOOL:       'tf',
          TType.BYTE:       'i8',
          TType.I16:        'i16',
          TType.I32:        'i32',
          TType.I64:        'i64',
          TType.DOUBLE:     'dbl',
          TType.STRING:     'str',
          TType.STRUCT:     'rec',
          TType.LIST:       'lst',
          TType.SET:        'set',
          TType.MAP:        'map'}

JTYPES = {}
for key in CTYPES.keys():
    JTYPES[CTYPES[key]] = key


class JSONBaseContext(object):

    def __init__(self, protocol):
        self.protocol = protocol
        self.first = True

    def doIO(self, function):
        pass

    def write(self):
        pass

    def read(self):
        pass

    def escapeNum(self):
        return False


class JSONListContext(JSONBaseContext):

    def doIO(self, function):
        if self.first is True:
            self.first = False
        else:
            function(COMMA)

    def write(self):
        self.doIO(self.protocol.trans.write)

    def read(self):
        self.doIO(self.protocol.readJSONSyntaxChar)


class JSONPairContext(JSONBaseContext):
    colon = True

    def doIO(self, function):
        if self.first is True:
            # Ignore extra commas at field start. Once context stack handling
            # fix is fully rolled out this can be removed.
            if self.protocol.reader.peek() == COMMA:
                self.protocol.readJSONSyntaxChar(COMMA)
            self.first = False
            self.colon = True
        else:
            function(COLON if self.colon is True else COMMA)
            self.colon = not self.colon

    def write(self):
        self.doIO(self.protocol.trans.write)

    def read(self):
        self.doIO(self.protocol.readJSONSyntaxChar)

    def escapeNum(self):
        return self.colon


class LookaheadReader():
    hasData = False
    data = ''

    def __init__(self, protocol):
        self.protocol = protocol

    def read(self):
        if self.hasData is True:
            self.hasData = False
        else:
            self.data = self.protocol.trans.read(1)
        if sys.version_info[0] >= 3 and isinstance(self.data, bytes):
            self.data = str(self.data, 'utf-8')
        return self.data

    def peek(self):
        if self.hasData is False:
            self.data = self.protocol.trans.read(1)
        self.hasData = True
        if sys.version_info[0] >= 3 and isinstance(self.data, bytes):
            self.data = str(self.data, 'utf-8')
        return self.data

class TJSONProtocolBase(TProtocolBase):

    def __init__(self, trans, validJSON=True):
        TProtocolBase.__init__(self, trans)
        self.validJSON = validJSON
        self.resetWriteContext()
        self.resetReadContext()

    def resetWriteContext(self):
        self.context = JSONBaseContext(self)
        self.contextStack = [self.context]

    def resetReadContext(self):
        self.resetWriteContext()
        self.reader = LookaheadReader(self)

    def pushContext(self, ctx):
        self.contextStack.append(ctx)
        self.context = ctx

    def popContext(self):
        self.contextStack.pop()
        if self.validJSON:
            self.context = self.contextStack[-1]

    def writeJSONString(self, string):
        # Python 3 JSON will not serialize bytes
        if isinstance(string, bytes) and sys.version_info.major >= 3:
            string = string.decode()
        self.context.write()
        self.trans.write(json.dumps(string))

    def writeJSONNumber(self, number):
        self.context.write()
        jsNumber = str(number)
        if self.context.escapeNum():
            jsNumber = "%s%s%s" % (QUOTE, jsNumber, QUOTE)
        self.trans.write(jsNumber)

    def writeJSONBase64(self, binary):
        self.context.write()
        self.trans.write(QUOTE)
        self.trans.write(base64.b64encode(binary))
        self.trans.write(QUOTE)

    def writeJSONObjectStart(self):
        self.context.write()
        self.trans.write(LBRACE)
        self.pushContext(JSONPairContext(self))

    def writeJSONObjectEnd(self):
        self.popContext()
        self.trans.write(RBRACE)

    def writeJSONArrayStart(self):
        self.context.write()
        self.trans.write(LBRACKET)
        self.pushContext(JSONListContext(self))

    def writeJSONArrayEnd(self):
        self.popContext()
        self.trans.write(RBRACKET)

    def readJSONSyntaxChar(self, character):
        current = self.reader.read()
        if character != current:
            raise TProtocolException(TProtocolException.INVALID_DATA,
                                     "Unexpected character: %s" % current)

    def readJSONString(self, skipContext):
        string = []
        if skipContext is False:
            self.context.read()
        self.readJSONSyntaxChar(QUOTE)
        while True:
            character = self.reader.read()
            if character == QUOTE:
                break
            if character == ESCSEQ[0]:
                character = self.reader.read()
                if character == ESCSEQ[1]:
                    self.readJSONSyntaxChar(ZERO)
                    self.readJSONSyntaxChar(ZERO)
                    data = self.trans.read(2)
                    if sys.version_info[0] >= 3 and isinstance(data, bytes):
                        character = json.JSONDecoder().decode(
                                '"\\u00%s"' % str(data, 'utf-8'))
                    else:
                        character = json.JSONDecoder().decode('"\\u00%s"' %
                                data)
                else:
                    off = ESCAPE_CHAR.find(character)
                    if off == -1:
                        raise TProtocolException(
                                TProtocolException.INVALID_DATA,
                                "Expected control char")
                    character = ESCAPE_CHAR_VALS[off]
            string.append(character)
        return ''.join(string)

    def isJSONNumeric(self, character):
        return (True if NUMERIC_CHAR.find(character) != - 1 else False)

    def readJSONQuotes(self):
        if (self.context.escapeNum()):
            self.readJSONSyntaxChar(QUOTE)

    def readJSONNumericChars(self):
        numeric = []
        while True:
            character = self.reader.peek()
            if self.isJSONNumeric(character) is False:
                break
            numeric.append(self.reader.read())
        return ''.join(numeric)

    def readJSONInteger(self):
        self.context.read()
        self.readJSONQuotes()
        numeric = self.readJSONNumericChars()
        self.readJSONQuotes()
        try:
            return int(numeric)
        except ValueError:
            raise TProtocolException(TProtocolException.INVALID_DATA,
                                     "Bad data encounted in numeric data")

    def readJSONDouble(self):
        self.context.read()
        if self.reader.peek() == QUOTE:
            string = self.readJSONString(True)
            try:
                double = float(string)
                if (self.context.escapeNum is False and
                    double != float('inf') and
                    double != float('-inf') and
                    double != float('nan')
                ):
                    raise TProtocolException(TProtocolException.INVALID_DATA,
                            "Numeric data unexpectedly quoted")
                return double
            except ValueError:
                raise TProtocolException(TProtocolException.INVALID_DATA,
                                         "Bad data encounted in numeric data")
        else:
            if self.context.escapeNum() is True:
                self.readJSONSyntaxChar(QUOTE)
            try:
                return float(self.readJSONNumericChars())
            except ValueError:
                raise TProtocolException(TProtocolException.INVALID_DATA,
                                         "Bad data encounted in numeric data")

    def readJSONBase64(self):
        string = self.readJSONString(False)
        return base64.b64decode(string)

    def readJSONObjectStart(self):
        self.context.read()
        self.readJSONSyntaxChar(LBRACE)
        self.pushContext(JSONPairContext(self))

    def readJSONObjectEnd(self):
        self.readJSONSyntaxChar(RBRACE)
        self.popContext()

    def readJSONArrayStart(self):
        self.context.read()
        self.readJSONSyntaxChar(LBRACKET)
        self.pushContext(JSONListContext(self))

    def readJSONArrayEnd(self):
        self.readJSONSyntaxChar(RBRACKET)
        self.popContext()


class TJSONProtocol(TJSONProtocolBase):

    def readMessageBegin(self):
        self.resetReadContext()
        self.readJSONArrayStart()
        if self.readJSONInteger() != VERSION:
            raise TProtocolException(TProtocolException.BAD_VERSION,
                                     "Message contained bad version.")
        name = self.readJSONString(False)
        typen = self.readJSONInteger()
        seqid = self.readJSONInteger()
        return (name, typen, seqid)

    def readMessageEnd(self):
        self.readJSONArrayEnd()

    def readStructBegin(self):
        self.readJSONObjectStart()

    def readStructEnd(self):
        self.readJSONObjectEnd()

    def readFieldBegin(self):
        character = self.reader.peek()
        ttype = 0
        id = 0
        if character == RBRACE:
            ttype = TType.STOP
        else:
            id = self.readJSONInteger()
            self.readJSONObjectStart()
            ttype = JTYPES[self.readJSONString(False)]
        return (None, ttype, id)

    def readFieldEnd(self):
        self.readJSONObjectEnd()

    def readMapBegin(self):
        self.readJSONArrayStart()
        keyType = JTYPES[self.readJSONString(False)]
        valueType = JTYPES[self.readJSONString(False)]
        size = self.readJSONInteger()
        self.readJSONObjectStart()
        return (keyType, valueType, size)

    def readMapEnd(self):
        self.readJSONObjectEnd()
        self.readJSONArrayEnd()

    def readCollectionBegin(self):
        self.readJSONArrayStart()
        elemType = JTYPES[self.readJSONString(False)]
        size = self.readJSONInteger()
        return (elemType, size)
    readListBegin = readCollectionBegin
    readSetBegin = readCollectionBegin

    def readCollectionEnd(self):
        self.readJSONArrayEnd()
    readSetEnd = readCollectionEnd
    readListEnd = readCollectionEnd

    def readBool(self):
        return (False if self.readJSONInteger() == 0 else True)

    def readNumber(self):
        return self.readJSONInteger()
    readByte = readNumber
    readI16 = readNumber
    readI32 = readNumber
    readI64 = readNumber

    def readDouble(self):
        return self.readJSONDouble()

    def readFloat(self):
        return self.readJSONDouble()

    def readString(self):
        string = self.readJSONString(False)
        if sys.version_info.major >= 3:
            # Generated code expects that protocols deal in bytes in Py3
            return string.encode('utf-8')
        return string

    def readBinary(self):
        return self.readJSONBase64()

    def writeMessageBegin(self, name, request_type, seqid):
        self.resetWriteContext()
        self.writeJSONArrayStart()
        self.writeJSONNumber(VERSION)
        self.writeJSONString(name)
        self.writeJSONNumber(request_type)
        self.writeJSONNumber(seqid)

    def writeMessageEnd(self):
        self.writeJSONArrayEnd()

    def writeStructBegin(self, name):
        self.writeJSONObjectStart()

    def writeStructEnd(self):
        self.writeJSONObjectEnd()

    def writeFieldBegin(self, name, ttype, id):
        self.writeJSONNumber(id)
        self.writeJSONObjectStart()
        self.writeJSONString(CTYPES[ttype])

    def writeFieldEnd(self):
        self.writeJSONObjectEnd()

    def writeFieldStop(self):
        pass

    def writeMapBegin(self, ktype, vtype, size):
        self.writeJSONArrayStart()
        self.writeJSONString(CTYPES[ktype])
        self.writeJSONString(CTYPES[vtype])
        self.writeJSONNumber(size)
        self.writeJSONObjectStart()

    def writeMapEnd(self):
        self.writeJSONObjectEnd()
        self.writeJSONArrayEnd()

    def writeListBegin(self, etype, size):
        self.writeJSONArrayStart()
        self.writeJSONString(CTYPES[etype])
        self.writeJSONNumber(size)

    def writeListEnd(self):
        self.writeJSONArrayEnd()

    def writeSetBegin(self, etype, size):
        self.writeJSONArrayStart()
        self.writeJSONString(CTYPES[etype])
        self.writeJSONNumber(size)

    def writeSetEnd(self):
        self.writeJSONArrayEnd()

    def writeBool(self, boolean):
        self.writeJSONNumber(1 if boolean is True else 0)

    def writeInteger(self, integer):
        self.writeJSONNumber(int(integer))
    writeByte = writeInteger
    writeI16 = writeInteger
    writeI32 = writeInteger
    writeI64 = writeInteger

    def writeDouble(self, dbl):
        self.writeJSONNumber(dbl)

    def writeFloat(self, flt):
        self.writeJSONNumber(flt)

    def writeString(self, string):
        self.writeJSONString(string)

    def writeBinary(self, binary):
        self.writeJSONBase64(binary)

class TJSONProtocolFactory:
    # validJSON specifies whether to emit valid JSON or possibly invalid but
    # backward-compatible one.
    def __init__(self, validJSON=True):
        # type: (bool) -> None
        self.validJSON = validJSON

    def getProtocol(self, trans):
        return TJSONProtocol(trans, self.validJSON)
