/*
 * 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.
 */

#include "TSimpleJSONProtocol.h"

using namespace apache::thrift::transport;
using namespace apache::thrift::reflection;

namespace apache { namespace thrift { namespace protocol {

TSimpleJSONProtocol::TSimpleJSONProtocol(std::shared_ptr<TTransport> ptrans):
                            TVirtualProtocol<TSimpleJSONProtocol,
                                             TJSONProtocol>(ptrans),
                                                 nextType_(nullptr),
                                                     numSkipped_(0) {
}

TSimpleJSONProtocol::TSimpleJSONProtocol(TTransport *ptrans):
                            TVirtualProtocol<TSimpleJSONProtocol,
                                             TJSONProtocol>(ptrans),
                                                 nextType_(nullptr),
                                                     numSkipped_(0) {
}


TSimpleJSONProtocol::~TSimpleJSONProtocol() {}

  /**
   * Writing functions.
   */

Schema * TSimpleJSONProtocol::getSchema() {
  return &schema_;
}

uint32_t TSimpleJSONProtocol::writeFieldBegin(const char* name,
                                        const TType fieldType,
                                        const int16_t fieldId) {
  return writeJSONString(name);
}


uint32_t TSimpleJSONProtocol::writeFieldEnd() {
  return 0;
}


uint32_t TSimpleJSONProtocol::writeMapBegin(const TType keyType,
                                      const TType valType,
                                      const uint32_t size) {
  return writeJSONObjectStart();
}


uint32_t TSimpleJSONProtocol::writeMapEnd() {
  return writeJSONObjectEnd();
}


uint32_t TSimpleJSONProtocol::writeListBegin(const TType elemType,
                                       const uint32_t size) {
  return writeJSONArrayStart();
}


uint32_t TSimpleJSONProtocol::writeSetBegin(const TType elemType,
                                      const uint32_t size) {
  return writeJSONArrayStart();
}


uint32_t TSimpleJSONProtocol::writeBool(const bool value) {
  return writeJSONBool(value);
}

  /**
   * Reading functions
   */


void TSimpleJSONProtocol::setNextStructType(uint64_t reflection_id) {
  nextType_ = getDataTypeFromTypeNum(reflection_id);
}

uint32_t TSimpleJSONProtocol::readStructBegin(std::string& name) {
  uint32_t result = TVirtualProtocol<TSimpleJSONProtocol,TJSONProtocol>::
                                    readStructBegin(name);
  enterType();
  return result;
}


uint32_t TSimpleJSONProtocol::readStructEnd() {
  uint32_t result = TVirtualProtocol<TSimpleJSONProtocol,TJSONProtocol>::
                                     readStructEnd();
  exitType();
  return result;
}


uint32_t TSimpleJSONProtocol::readFieldBegin(std::string& name,
                                    TType& fieldType,
                                    int16_t& fieldId) {
  uint32_t result = 0;

  auto currentType = getCurrentDataType();

  skipWhitespace();
  result += getNumSkippedChars();

  // Check if we hit the end of the list
  uint8_t ch = reader_.peek();
  if (ch == kJSONObjectEnd) {
    fieldType = T_STOP;
    return result;
  }

  std::string tmpStr;
  result += readJSONString(tmpStr);

  if (currentType != nullptr) {
    if (! currentType->__isset.fields) {
      throw TProtocolException(TProtocolException::INVALID_DATA,
                                   "Expected a struct type, but "
                                   "actually not a struct");
    }

    auto& fields = currentType->fields;

    // find the corresponding StructField object and field id of the field
    for (auto ite = fields.begin(); ite != fields.end(); ++ite) {
      if (ite->second.name == tmpStr) {
        fieldId = ite->first;
        fieldType = getTypeIdFromTypeNum(ite->second.type);

        // set the nextType_ if the field type is a compound type
        // e.g. list<i64>, mySimpleStruct
        //
        // this is not really necessary, because before calling
        // readStructBegin(), setNextStructType() should be called
        // but this allows only calling setNextStructType() on the
        // base type
        auto& field = ite->second;
        if (isCompoundType(field.type)) {
          nextType_ = getDataTypeFromTypeNum(field.type);
        }

        return result;
      }
    }
  }

  // if the field is not found or
  // the entire struct is being skipped

  fieldId = 0;

  skipWhitespace();
  uint8_t delimiter = reader_.read(); // delimiter should be ':'

  fieldType = guessTypeIdFromFirstByte();

  bool wasPutBack = reader_.put(delimiter);

  assert(wasPutBack);

  return result + getNumSkippedChars();
}


uint32_t TSimpleJSONProtocol::readFieldEnd() {
  return 0;
}

uint32_t TSimpleJSONProtocol::readMapBegin(TType& keyType,
                                    TType& valType,
                                    uint32_t& size,
                                    bool& sizeUnknown) {
  enterType();

  auto currentType = getCurrentDataType();
  bool beingSkipped = (currentType == nullptr);

  // since we never guess an unknown field to have a map type
  // we should never arrive here
  assert(!beingSkipped);

  keyType = getTypeIdFromTypeNum(currentType->mapKeyType);
  valType = getTypeIdFromTypeNum(currentType->valueType);
  size = 0;
  sizeUnknown = true;

  if (isCompoundType(currentType->mapKeyType)) {
    nextType_ = getDataTypeFromTypeNum(currentType->mapKeyType);

  } else if (isCompoundType(currentType->valueType)) {
    nextType_ = getDataTypeFromTypeNum(currentType->valueType);
  }

  return readJSONObjectStart();
}


bool TSimpleJSONProtocol::peekMap() {
  skipWhitespace();
  return reader_.peek() != kJSONObjectEnd;
}


uint32_t TSimpleJSONProtocol::readMapEnd() {
  uint32_t result = getNumSkippedChars() + readJSONObjectEnd();
  exitType();
  return result;
}


uint32_t TSimpleJSONProtocol::readListBegin(TType& elemType,
                                   uint32_t& size,
                                   bool& sizeUnknown) {
  enterType();

  auto currentType = getCurrentDataType();
  bool beingSkipped = (currentType == nullptr);

  if (beingSkipped) {
    uint32_t result = readJSONArrayStart();
    elemType = guessTypeIdFromFirstByte();
    size = 0;
    sizeUnknown = true;
    return result + getNumSkippedChars();

  } else {
    elemType = getTypeIdFromTypeNum(currentType->valueType);
    size = 0;
    sizeUnknown = true;

    if (isCompoundType(currentType->valueType)) {
      nextType_ = getDataTypeFromTypeNum(currentType->valueType);
    }

    return readJSONArrayStart();

  }
}


bool TSimpleJSONProtocol::peekList() {
  skipWhitespace();
  return reader_.peek() != kJSONArrayEnd;
}


uint32_t TSimpleJSONProtocol::readListEnd() {
  uint32_t result = getNumSkippedChars();
  result += TVirtualProtocol<TSimpleJSONProtocol,TJSONProtocol>::
                            readListEnd();
  exitType();
  return result;
}


uint32_t TSimpleJSONProtocol::readSetBegin(TType& elemType,
                                  uint32_t& size,
                                  bool& sizeUnknown) {
  enterType();

  auto currentType = getCurrentDataType();
  bool beingSkipped = (currentType == nullptr);

  // since we never guess an unknown field to have a set type
  // we should never arrive here
  assert(!beingSkipped);

  elemType = getTypeIdFromTypeNum(currentType->valueType);
  size = 0;
  sizeUnknown = true;

  if (isCompoundType(currentType->valueType)) {
    nextType_ = getDataTypeFromTypeNum(currentType->valueType);
  }

  return readJSONArrayStart();
}


bool TSimpleJSONProtocol::peekSet() {
  skipWhitespace();
  return reader_.peek() != kJSONArrayEnd;
}


uint32_t TSimpleJSONProtocol::readSetEnd() {
  uint32_t result = getNumSkippedChars();
  result += TVirtualProtocol<TSimpleJSONProtocol,TJSONProtocol>::
                             readSetEnd();
  exitType();
  return result;
}

uint32_t TSimpleJSONProtocol::readBool(bool& value) {
  return readJSONBool(value);
}

TType TSimpleJSONProtocol::getTypeIdFromTypeNum(int64_t fieldType) {
  Type type = getType(fieldType);
  switch (type) {
    case TYPE_VOID:   return T_VOID;
    case TYPE_STRING: return T_STRING;
    case TYPE_BOOL:   return T_BOOL;
    case TYPE_BYTE:   return T_BYTE;
    case TYPE_I16:    return T_I16;
    case TYPE_I32:    return T_I32;
    case TYPE_I64:    return T_I64;
    case TYPE_DOUBLE: return T_DOUBLE;
    case TYPE_FLOAT:  return T_FLOAT;
    case TYPE_LIST:   return T_LIST;
    case TYPE_SET:    return T_SET;
    case TYPE_MAP:    return T_MAP;
    case TYPE_STRUCT: return T_STRUCT;
    case TYPE_ENUM:   return T_I32;
    case TYPE_SERVICE:
    case TYPE_PROGRAM:
    default:
      throw TProtocolException(TProtocolException::NOT_IMPLEMENTED,
                               "Unrecognized type");
  }
}

/**
 * Given a byte, peeked from the beginning of some JSON value, determine a type
 * of that value. Result type is used to decide how to skip that value, so
 * differences between compatible types don't matter and the more general type
 * is assumed.
 * STOP is returned for } and ] to indicate end of collection.
 */
TType TSimpleJSONProtocol::guessTypeIdFromFirstByte() {
  skipWhitespace();
  uint8_t byte = reader_.peek();

  if (byte == kJSONObjectEnd ||
      byte == kJSONArrayEnd) {
    return T_STOP;
  } else if (byte == kJSONStringDelimiter) {
    return T_STRING;
  } else if (byte == kJSONObjectStart) {
    return T_STRUCT;
  } else if (byte == kJSONArrayStart) {
    return T_LIST;
  } else if (byte == kJSONTrue[0] ||
             byte == kJSONFalse[0]) {
    return T_BOOL;
  } else if (byte == '+' ||
             byte == '-' ||
             byte == '.' ||
             byte == '0' ||
             byte == '1' ||
             byte == '2' ||
             byte == '3' ||
             byte == '4' ||
             byte == '5' ||
             byte == '6' ||
             byte == '7' ||
             byte == '8' ||
             byte == '9') {
        return T_DOUBLE;
  } else {
    throw TProtocolException(TProtocolException::NOT_IMPLEMENTED,
                "Unrecognized byte: " + std::string((char *)&byte, 1));
  }
}

DataType* TSimpleJSONProtocol::getDataTypeFromTypeNum(int64_t typeNum) {
  auto ite = schema_.dataTypes.find(typeNum);

  if (ite == schema_.dataTypes.cend()) {
    throw TProtocolException(TProtocolException::INVALID_DATA,
                             "Type id not found, schema is corrupted");
  }

  return &(ite->second);
}


void TSimpleJSONProtocol::enterType() {
  typeStack_.push(nextType_);
  nextType_ = nullptr;
}


void TSimpleJSONProtocol::exitType() {
  auto lastType = getCurrentDataType();
  typeStack_.pop();

  auto currentType = getCurrentDataType();

  if (currentType == nullptr) {
    nextType_ = nullptr;

  } else if (currentType->__isset.mapKeyType &&
             isCompoundType(currentType->mapKeyType) &&
             (!currentType->__isset.valueType ||
              !isCompoundType(currentType->valueType) ||
              lastType == getDataTypeFromTypeNum(currentType->valueType))) {
    nextType_ = getDataTypeFromTypeNum(currentType->mapKeyType);

  } else if (currentType->__isset.valueType &&
             isCompoundType(currentType->valueType)) {
    nextType_ = getDataTypeFromTypeNum(currentType->valueType);

  } else {
    nextType_ = nullptr;
  }
}

// returns nullptr when the struct is to be skipped
const DataType * TSimpleJSONProtocol::getCurrentDataType() {
  if (!typeStack_.empty()) {
    return typeStack_.top();
  } else {
    return nullptr;
  }
}

void TSimpleJSONProtocol::skipWhitespace() {
  numSkipped_ += this->skipJSONWhitespace();
}

uint32_t TSimpleJSONProtocol::getNumSkippedChars() {
  auto temp = numSkipped_;
  numSkipped_ = 0;
  return temp;
}

bool TSimpleJSONProtocol::isCompoundType(int64_t fieldType) {
  return !isBaseType(getType(fieldType));
}

}}} // apache::thrift::protocol
