/*
 * Copyright 2015 Facebook, Inc.
 *
 * 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.
 */
#define __STDC_LIMIT_MACROS
#include <Python.h>

#include <stdint.h>
#include <type_traits>

#include <thrift/lib/cpp2/protocol/BinaryProtocol.h>
#include <thrift/lib/cpp2/protocol/CompactProtocol.h>
#include <thrift/lib/cpp2/protocol/ProtocolReaderWithRefill.h>

#include <folly/Range.h>
#include <folly/ScopeGuard.h>

using apache::thrift::protocol::TType;
using apache::thrift::BinaryProtocolReaderWithRefill;
using apache::thrift::BinaryProtocolWriter;
using apache::thrift::CompactProtocolReaderWithRefill;
using apache::thrift::CompactProtocolWriter;

#if PY_MAJOR_VERSION >= 3
  #define FROM_LONG PyLong_FromLong
  #define AS_LONG PyLong_AsLong
#else
  #define FROM_LONG PyInt_FromLong
  #define AS_LONG PyInt_AsLong
  #include <cStringIO.h>
#endif

#if PY_MAJOR_VERSION >= 3
static PyTypeObject* BytesIOType;
#endif

// Stolen from cStringIO.c and also works for Python 3
typedef struct {
  PyObject_HEAD
  char *buf;
  Py_ssize_t pos, string_size;
} IOobject;

#define IOOOBJECT(O) ((IOobject*)(O))

// Stolen from fastbinary.c
#define INTERN_STRING(value) _intern_ ## value

static PyObject* INTERN_STRING(cstringio_buf);
static PyObject* INTERN_STRING(cstringio_refill);

typedef struct {
  PyObject* stringiobuf;
  PyObject* refill_callable;
} DecodeBuffer;

static void
free_decodebuf(DecodeBuffer* d) {
  Py_XDECREF(d->stringiobuf);
  Py_XDECREF(d->refill_callable);
}

static bool
decode_buffer_from_obj(DecodeBuffer* dest, PyObject* obj) {
  dest->stringiobuf = PyObject_GetAttr(obj, INTERN_STRING(cstringio_buf));
  if (!dest->stringiobuf) {
    return false;
  }

#if PY_MAJOR_VERSION >= 3
  if (!PyObject_TypeCheck(dest->stringiobuf, BytesIOType)) {
#else
  if (!PycStringIO_InputCheck(dest->stringiobuf)) {
#endif
    free_decodebuf(dest);
    PyErr_SetString(PyExc_TypeError, "expecting stringio input");
    return false;
  }

  dest->refill_callable = PyObject_GetAttr(obj,
      INTERN_STRING(cstringio_refill));
  if(!dest->refill_callable) {
    free_decodebuf(dest);
    return false;
  }

  if (!PyCallable_Check(dest->refill_callable)) {
    free_decodebuf(dest);
    PyErr_SetString(PyExc_TypeError, "expecting callable");
    return false;
  }

  return true;
}

#define INT_CONV_ERROR_OCCURRED(v) ( ((v) == -1) && PyErr_Occurred() )
#define CHECK_RANGE(v, min, max) ( ((v) <= (max)) && ((v) >= (min)) )

typedef struct {
  TType element_type;
  PyObject* typeargs;
} SetListTypeArgs;

typedef struct {
  TType ktype;
  TType vtype;
  PyObject* ktypeargs;
  PyObject* vtypeargs;
} MapTypeArgs;

typedef struct {
  PyObject* klass;
  PyObject* spec;
  bool isunion;
} StructTypeArgs;

typedef struct {
  int tag;
  TType type;
  PyObject* attrname;
  PyObject* typeargs;
  PyObject* defval;
} StructItemSpec;

static bool
parse_set_list_args(SetListTypeArgs* dest, PyObject* typeargs) {
  long element_type;

  if (PyTuple_Size(typeargs) != 2) {
    PyErr_SetString(PyExc_TypeError,
        "expecting tuple of size 2 for list/set type args");
    return false;
  }

  element_type = AS_LONG(PyTuple_GET_ITEM(typeargs, 0));
  if (INT_CONV_ERROR_OCCURRED(element_type)) {
    return false;
  }

  dest->element_type = static_cast<TType>(element_type);
  dest->typeargs = PyTuple_GET_ITEM(typeargs, 1);

  return true;
}

static bool
parse_map_args(MapTypeArgs* dest, PyObject* typeargs) {
  long ktype, vtype;

  if (PyTuple_Size(typeargs) != 4) {
    PyErr_SetString(PyExc_TypeError,
        "expecting tuple of size 4 for map type args");
    return false;
  }

  ktype = AS_LONG(PyTuple_GET_ITEM(typeargs, 0));
  if (INT_CONV_ERROR_OCCURRED(ktype)) {
    return false;
  }

  vtype = AS_LONG(PyTuple_GET_ITEM(typeargs, 2));
  if (INT_CONV_ERROR_OCCURRED(vtype)) {
    return false;
  }

  dest->ktype = static_cast<TType>(ktype);
  dest->vtype = static_cast<TType>(vtype);
  dest->ktypeargs = PyTuple_GET_ITEM(typeargs, 1);
  dest->vtypeargs = PyTuple_GET_ITEM(typeargs, 3);

  return true;
}

static bool
parse_struct_args(StructTypeArgs* dest, PyObject* typeargs) {
  if (PyList_Size(typeargs) != 3) {
    PyErr_SetString(PyExc_TypeError,
        "expecting list of size 3 for struct args");
    return false;
  }

  dest->klass = PyList_GET_ITEM(typeargs, 0);
  dest->spec = PyList_GET_ITEM(typeargs, 1);
  dest->isunion = (PyObject_IsTrue(PyList_GET_ITEM(typeargs, 2)) == 1);

  return true;
}

static int
parse_struct_item_spec(StructItemSpec* dest, PyObject* spec_tuple) {
  long tag, type;

  // i'd like to use ParseArgs here, but it seems to be a bottleneck.
  if (PyTuple_Size(spec_tuple) != 6) {
    PyErr_SetString(PyExc_TypeError, "expecting 6 arguments for spec tuple");
    return false;
  }

  tag = AS_LONG(PyTuple_GET_ITEM(spec_tuple, 0));
  if (INT_CONV_ERROR_OCCURRED(tag)) {
    return false;
  }

  type = AS_LONG(PyTuple_GET_ITEM(spec_tuple, 1));
  if (INT_CONV_ERROR_OCCURRED(type)) {
    return false;
  }

  dest->tag = tag;
  dest->type = static_cast<TType>(type);
  dest->attrname = PyTuple_GET_ITEM(spec_tuple, 2);
  dest->typeargs = PyTuple_GET_ITEM(spec_tuple, 3);
  dest->defval = PyTuple_GET_ITEM(spec_tuple, 4);
  // Arg #5 is the 'required' field, which is ignored when serializing.

  return true;
}

static inline bool
check_ssize_t_32(Py_ssize_t len) {
  // error from getting the int
  if (INT_CONV_ERROR_OCCURRED(len)) {
    return false;
  }
  if (!CHECK_RANGE(len, 0, INT32_MAX)) {
    PyErr_SetString(PyExc_OverflowError, "string size out of range");
    return false;
  }
  return true;
}

static inline bool
parse_pyint(PyObject* o, int32_t* ret, int32_t min, int32_t max) {
  long val = AS_LONG(o);
  if (INT_CONV_ERROR_OCCURRED(val)) {
    return false;
  }
  if (!CHECK_RANGE(val, min, max)) {
    PyErr_SetString(PyExc_OverflowError, "int out of range");
    return false;
  }

  *ret = (int32_t) val;
  return true;
}

static inline bool
parse_pyfloat(PyObject *o, double *ret) {
  double val = PyFloat_AsDouble(o);
  if (val == -1.0 && PyErr_Occurred()) {
    return false;
  }
  *ret = val;
  return true;
}

/**
 * Main loop for reading a struct and serializing fields.
 */
template<typename Writer>
static bool encode_impl(Writer *writer, PyObject *value, PyObject *typeargs,
    TType type, int utf8strings) {
  switch (type) {
    case TType::T_BOOL: {
      int v = PyObject_IsTrue(value);
      if (v == -1) {
        return false;
      }
      writer->writeBool(v == 1);
      break;
    }
    case TType::T_I08: {
      int32_t val;
      if (!parse_pyint(value, &val, INT8_MIN, INT8_MAX)) {
        return false;
      }
      writer->writeByte((int8_t)val);
      break;
    }
    case TType::T_I16: {
      int32_t val;
      if (!parse_pyint(value, &val, INT16_MIN, INT16_MAX)) {
        return false;
      }
      writer->writeI16((int16_t)val);
      break;
    }
    case TType::T_I32: {
      int32_t val;
      if (!parse_pyint(value, &val, INT32_MIN, INT32_MAX)) {
        return false;
      }
      writer->writeI32(val);
      break;
    }
    case TType::T_I64: {
      int64_t val = PyLong_AsLongLong(value);
      if (INT_CONV_ERROR_OCCURRED(val)) {
        return false;
      }
      if (!CHECK_RANGE(val, INT64_MIN, INT64_MAX)) {
        PyErr_SetString(PyExc_OverflowError, "int out of range");
        return false;
      }
      writer->writeI64(val);
      break;
    }
    case TType::T_DOUBLE: {
      double val;
      if (!parse_pyfloat(value, &val)) {
        return false;
      }
      writer->writeDouble(val);
      break;
    }
    case TType::T_FLOAT: {
      double val;
      if (!parse_pyfloat(value, &val)) {
        return false;
      }
      writer->writeFloat((float)val);
      break;
    }
    case TType::T_STRING: {
      bool encoded = false;
      Py_ssize_t len;

#if PY_MAJOR_VERSION >= 3
      if (!PyBytes_Check(value)) {
        value = PyUnicode_AsUTF8String(value);
        if (value == nullptr) {
          PyErr_SetString(PyExc_TypeError, "can not encode");
          return false;
        }
        encoded = true;
      }
      len = PyBytes_Size(value);
#else
      if (utf8strings && value->ob_type == &PyUnicode_Type) {
        value = PyUnicode_AsUTF8String(value);
        if (value == nullptr) {
          PyErr_SetString(PyExc_TypeError, "can not encode using utf8");
          return false;
        }
        encoded = true;
      }
      len = PyString_Size(value);
#endif

      if (!check_ssize_t_32(len)) {
        return false;
      }

#if PY_MAJOR_VERSION >= 3
      folly::StringPiece str(PyBytes_AsString(value), len);
#else
      folly::StringPiece str(PyString_AsString(value), len);
#endif
      writer->writeString(str);
      if (encoded) {
        Py_DECREF(value);
      }
      break;
    }
    case TType::T_LIST:
    case TType::T_SET: {
      Py_ssize_t len;
      SetListTypeArgs parsedargs;
      PyObject *item;
      PyObject *iterator;

      if (!parse_set_list_args(&parsedargs, typeargs)) {
        return false;
      }

      len = PyObject_Length(value);

      if (!check_ssize_t_32(len)) {
        return false;
      }

      writer->writeListBegin(parsedargs.element_type, (uint32_t)len);
      iterator = PyObject_GetIter(value);
      if (iterator == nullptr) {
        return false;
      }

      while ((item = PyIter_Next(iterator))) {
        if (!encode_impl(writer, item, parsedargs.typeargs,
              parsedargs.element_type, utf8strings)) {
          Py_DECREF(item);
          Py_DECREF(iterator);
          return false;
        }
        Py_DECREF(item);
      }

      Py_DECREF(iterator);
      writer->writeListEnd();

      if (PyErr_Occurred()) {
        return false;
      }
      break;
    }
    case TType::T_MAP: {
      PyObject *k, *v;
      Py_ssize_t pos = 0, len;

      MapTypeArgs parsedargs;

      len = PyDict_Size(value);
      if (!check_ssize_t_32(len)) {
        return false;
      }

      if (!parse_map_args(&parsedargs, typeargs)) {
        return false;
      }

      writer->writeMapBegin(parsedargs.ktype, parsedargs.vtype, (uint32_t)len);
      while (PyDict_Next(value, &pos, &k, &v)) {
        Py_INCREF(k);
        Py_INCREF(v);

        if (!encode_impl(writer, k, parsedargs.ktypeargs, parsedargs.ktype,
              utf8strings)
            || !encode_impl(writer, v, parsedargs.vtypeargs, parsedargs.vtype,
              utf8strings)) {
          Py_DECREF(k);
          Py_DECREF(v);
          return false;
        }
        Py_DECREF(k);
        Py_DECREF(v);
      }
      writer->writeMapEnd();
      break;
    }
    case TType::T_STRUCT: {
      StructTypeArgs parsedargs;

      if (!parse_struct_args(&parsedargs, typeargs)) {
        return false;
      }

      Py_ssize_t nspec = PyTuple_Size(parsedargs.spec);
      if (nspec == -1) {
        return false;
      }

      if (parsedargs.isunion) {
        // Union only has a field and a value.
        writer->writeStructBegin("");
        PyObject *field = PyObject_GetAttrString(value, "field");
        if (!field) {
          return false;
        }

        int fid = static_cast<int>(AS_LONG(field));
        PyObject *spec_tuple = PyTuple_GET_ITEM(parsedargs.spec, fid);
        if (spec_tuple != Py_None) {
          StructItemSpec parsedspec;
          PyObject* instval = nullptr;

          if (!parse_struct_item_spec(&parsedspec, spec_tuple)) {
            return false;
          }

          instval = PyObject_GetAttrString(value, "value");
          if (!instval) {
            return false;
          }

          if (instval == Py_None) {
            Py_DECREF(instval);
            return false;
          }

#if PY_MAJOR_VERSION >= 3
          const char* fieldname = PyUnicode_AsUTF8(parsedspec.attrname);
#else
          const char* fieldname = PyString_AsString(parsedspec.attrname);
#endif
          writer->writeFieldBegin(fieldname, parsedspec.type, parsedspec.tag);
          if (!encode_impl(writer, instval, parsedspec.typeargs,
                parsedspec.type, utf8strings)) {
            Py_DECREF(instval);
            return false;
          }

          Py_DECREF(instval);
          writer->writeFieldEnd();
        }

        writer->writeFieldStop();
        writer->writeStructEnd();
        break;
      }

      // Both binary and compact ignore the struct name.
      writer->writeStructBegin("");
      for (Py_ssize_t i = 0; i < nspec; i++) {
        StructItemSpec parsedspec;
        PyObject* spec_tuple;
        PyObject* instval = nullptr;

        spec_tuple = PyTuple_GET_ITEM(parsedargs.spec, i);
        if (spec_tuple == Py_None) {
          continue;
        }

        if (!parse_struct_item_spec(&parsedspec, spec_tuple)) {
          return false;
        }

        instval = PyObject_GetAttr(value, parsedspec.attrname);
        if (!instval) {
          return false;
        }

        if (instval == Py_None) {
          Py_DECREF(instval);
          continue;
        }

#if PY_MAJOR_VERSION >= 3
        const char* fieldname = PyUnicode_AsUTF8(parsedspec.attrname);
#else
        const char* fieldname = PyString_AsString(parsedspec.attrname);
#endif
        writer->writeFieldBegin(fieldname, parsedspec.type, parsedspec.tag);
        if (!encode_impl(writer, instval, parsedspec.typeargs, parsedspec.type,
              utf8strings)) {
          Py_DECREF(instval);
          return false;
        }

        Py_DECREF(instval);
        writer->writeFieldEnd();
      }

      writer->writeFieldStop();
      writer->writeStructEnd();
      break;
    }
    case TType::T_STOP:
    case TType::T_VOID:
    case TType::T_UTF8:
    case TType::T_UTF16:
    case TType::T_U64:
    case TType::T_STREAM:
    default:
      PyErr_SetString(PyExc_TypeError, "Unexpected TType");
      return false;
  }

  return true;
}

template<typename Reader>
static PyObject*
decode_val(Reader *reader, TType type, PyObject *typeargs, int utf8strings);

template<typename Reader>
static bool
decode_struct(Reader *reader, PyObject *value, StructTypeArgs *args,
    int utf8strings) {
  int speclen = PyTuple_Size(args->spec);
  if (speclen == -1) {
    return false;
  }

  std::string sname;
  reader->readStructBegin(sname);

  std::string fname;
  TType ftype;
  int16_t fid;
  PyObject *itemspec;
  StructItemSpec parsedspec;

  int first_tag = 0;
  bool first_tag_read = false;
  PyObject *first_item_spec;
  StructItemSpec first_parsed_spec;

  while (true) {
    PyObject *fieldval = nullptr;
    reader->readFieldBegin(fname, ftype, fid);
    if (ftype == TType::T_STOP) {
      break;
    }

    if (!first_tag_read) {
      first_tag_read = true;
      if (speclen == 0) { // Empty struct and all fields will be skipped
        first_item_spec = Py_None;
      } else {
        first_item_spec = PyTuple_GET_ITEM(args->spec, 0);
        if (first_item_spec != Py_None) {
          if (!parse_struct_item_spec(&first_parsed_spec, first_item_spec)) {
            return false;
          }
          first_tag = first_parsed_spec.tag;
        }
      }
    }

    fid -= first_tag;
    if (fid > 0 && fid < speclen) {
      itemspec = PyTuple_GET_ITEM(args->spec, fid);
    } else if (fid == 0) {
      itemspec = first_item_spec;
    } else {
      itemspec = Py_None;
    }

    if (itemspec == Py_None) {
      if (!reader->skip(ftype)) {
        if (ftype == TType::T_BOOL &&
            std::is_same<Reader, CompactProtocolReaderWithRefill>::value) {
          // readBool may return 0 in compact.
          continue;
        }
        PyErr_SetString(PyExc_TypeError,
            "Can't skip nonexistent struct field");
        return false;
      } else {
        continue;
      }
    }

    if (fid == 0) {
      parsedspec = first_parsed_spec;
    } else if (!parse_struct_item_spec(&parsedspec, itemspec)) {
      return false;
    }

    if (parsedspec.type != ftype) {
      if (!reader->skip(ftype)) {
        if (ftype == TType::T_BOOL &&
            std::is_same<Reader, CompactProtocolReaderWithRefill>::value) {
          // readBool may return 0 in compact.
          continue;
        }
        PyErr_SetString(PyExc_TypeError,
            "struct field had wrong type while reading and can't be skipped");
        return false;
      } else {
        continue;
      }
    }

    fieldval = decode_val(reader, parsedspec.type, parsedspec.typeargs,
        utf8strings);
    if (!fieldval) {
      return false;
    }

    if (args->isunion){
      PyObject *fieldobj = PyObject_GetAttrString(value, "field");
      if (!fieldobj) {
        Py_DECREF(fieldval);
        return false;
      }
      int field = (int)AS_LONG(fieldobj);
      if (field == -1 && PyErr_Occurred()) {
        Py_DECREF(fieldval);
        Py_DECREF(fieldobj);
        return false;
      }
      if (field != 0) {
        PyErr_SetString(PyExc_AssertionError, "field already set in union");
        Py_DECREF(fieldval);
        Py_DECREF(fieldobj);
        return false;
      }
      Py_DECREF(fieldobj);

      PyObject *valueobj = PyObject_GetAttrString(value, "value");
      if (!valueobj) {
        Py_DECREF(fieldobj);
        return false;
      }
      if (valueobj != Py_None) {
        PyErr_SetString(PyExc_AssertionError, "value already set in union");
        Py_DECREF(fieldval);
        Py_DECREF(valueobj);
        return false;
      }
      Py_DECREF(valueobj);

      PyObject *tagobj = FROM_LONG(parsedspec.tag);
      if (!tagobj) {
        return false;
      }
      if (PyObject_SetAttrString(value, "value", fieldval) == -1 ||
          PyObject_SetAttrString(value, "field", tagobj) == -1) {
        Py_DECREF(fieldval);
        Py_DECREF(tagobj);
        return false;
      }
      Py_DECREF(tagobj);
    } else {
      if (PyObject_SetAttr(value, parsedspec.attrname, fieldval) == -1) {
        Py_DECREF(fieldval);
        return false;
      }
    }
    Py_DECREF(fieldval);
    reader->readFieldEnd();
  }

  reader->readStructEnd();
  return true;
}

template<typename Reader>
static PyObject*
decode_val(Reader *reader, TType type, PyObject *typeargs, int utf8strings) {
  switch (type) {
    case TType::T_BOOL: {
      bool v;
      reader->readBool(v);
      if (v) {
        Py_RETURN_TRUE;
      } else {
        Py_RETURN_FALSE;
      }
    }
    case TType::T_I08: {
      int8_t v;
      reader->readByte(v);
      return FROM_LONG(v);
    }
    case TType::T_I16: {
      int16_t v;
      reader->readI16(v);
      return FROM_LONG(v);
    }
    case TType::T_I32: {
      int32_t v;
      reader->readI32(v);
      return FROM_LONG(v);
    }
    case TType::T_I64: {
      int64_t v;
      reader->readI64(v);
      if (CHECK_RANGE(v, LONG_MIN, LONG_MAX)) {
        return FROM_LONG((long) v);
      }
      return PyLong_FromLongLong(v);
    }
    case TType::T_DOUBLE: {
      double v;
      reader->readDouble(v);
      return PyFloat_FromDouble(v);
    }
    case TType::T_FLOAT: {
      float v;
      reader->readFloat(v);
      return PyFloat_FromDouble((double) v);
    }
    case TType::T_STRING: {
      std::string s;
      reader->readString(s);
      if (utf8strings && PyObject_IsTrue(typeargs)) {
        return PyUnicode_FromStringAndSize(s.data(), s.length());
      } else {
#if PY_MAJOR_VERSION >= 3
        return PyBytes_FromStringAndSize(s.data(), s.length());
#else
        return PyString_FromStringAndSize(s.data(), s.length());
#endif
      }
    }
    case TType::T_LIST:
    case TType::T_SET: {
      SetListTypeArgs parsedargs;
      TType ttype;
      uint32_t len;

      if (!parse_set_list_args(&parsedargs, typeargs)) {
        return nullptr;
      }

      reader->readListBegin(ttype, len);
      if (ttype != parsedargs.element_type) {
        PyErr_SetString(PyExc_TypeError,
            "got wrong ttype while reading list/set field");
        return nullptr;
      }

      PyObject *ret = PyList_New(len);
      if (!ret) {
        return nullptr;
      }

      for (int i = 0; i < len; i++) {
        PyObject *item = decode_val(reader, parsedargs.element_type,
            parsedargs.typeargs, utf8strings);
        if (!item) {
          Py_DECREF(ret);
          return nullptr;
        }
        PyList_SET_ITEM(ret, i, item);
      }

      reader->readListEnd();
      if (type == TType::T_SET) {
        PyObject *setret = PySet_New(ret);
        Py_DECREF(ret);
        return setret;
      }
      return ret;
    }
    case TType::T_MAP: {
      MapTypeArgs parsedargs;
      TType ktype, vtype;
      uint32_t len;

      if (!parse_map_args(&parsedargs, typeargs)) {
        return nullptr;
      }

      reader->readMapBegin(ktype, vtype, len);
      if (ktype != parsedargs.ktype || vtype != parsedargs.vtype) {
        if (len == 0 &&
            std::is_same<Reader, CompactProtocolReaderWithRefill>::value) {
          reader->readMapEnd();
          return PyDict_New();
        }
        PyErr_SetString(PyExc_TypeError,
            "got wrong ttype while reading map field");
        return nullptr;
      }

      PyObject *ret = PyDict_New();
      if (!ret) {
        return nullptr;
      }

      for (int i = 0; i < len; i++) {
        PyObject *k = nullptr;
        PyObject *v = nullptr;
        k = decode_val(reader, parsedargs.ktype, parsedargs.ktypeargs,
            utf8strings);
        if (!k) {
          Py_DECREF(ret);
          return nullptr;
        }
        v = decode_val(reader, parsedargs.vtype, parsedargs.vtypeargs,
            utf8strings);
        if (!v) {
          Py_DECREF(ret);
          Py_DECREF(k);
          return nullptr;
        }
        if (PyDict_SetItem(ret, k, v) == -1) {
          Py_DECREF(ret);
          Py_DECREF(k);
          Py_DECREF(v);
          return nullptr;
        }

        Py_DECREF(k);
        Py_DECREF(v);
      }

      reader->readMapEnd();
      return ret;
    }
    case TType::T_STRUCT: {
      StructTypeArgs parsedargs;
      PyObject *ret;

      if (!parse_struct_args(&parsedargs, typeargs)) {
        return nullptr;
      }

      ret = PyObject_CallObject(parsedargs.klass, nullptr);
      if (!ret) {
        return nullptr;
      }

      if (!decode_struct(reader, ret, &parsedargs, utf8strings)) {
        Py_DECREF(ret);
        return nullptr;
      }

      return ret;
    }
    case TType::T_STOP:
    case TType::T_VOID:
    case TType::T_UTF16:
    case TType::T_UTF8:
    case TType::T_U64:
    default:
      PyErr_SetString(PyExc_TypeError, "Unexpected TType");
      return nullptr;
  }
}

/* --- TOP-LEVEL WRAPPER AND TEMPLATES --- */

template<typename Writer>
static PyObject*
encodeT(PyObject *enc_obj, PyObject *spec, int utf8strings) {
  folly::IOBufQueue encoded(folly::IOBufQueue::cacheChainLength());
  Writer writer;
  writer.setOutput(&encoded);

  encode_impl(&writer, enc_obj, spec, TType::T_STRUCT, utf8strings);

  if (PyErr_Occurred()) {
    return nullptr;
  }

  auto buf = encoded.split(encoded.chainLength());
  buf->coalesce();
#if PY_MAJOR_VERSION >= 3
  return PyBytes_FromStringAndSize((const char*)buf->data(), buf->length());
#else
  return PyString_FromStringAndSize((const char*)buf->data(), buf->length());
#endif
}

static std::unique_ptr<folly::IOBuf> refill(DecodeBuffer *input,
    const uint8_t *remaining, int rlen, int read, int len) {
  PyObject *newiobuf;

  IOobject *ioobj = IOOOBJECT(input->stringiobuf);
  ioobj->pos += read;

#if PY_MAJOR_VERSION >= 3
  newiobuf = PyObject_CallFunction(input->refill_callable, (char*)"y#i",
      remaining, rlen, len, nullptr);
#else
  newiobuf = PyObject_CallFunction(input->refill_callable, (char*)"s#i",
      remaining, rlen, len, nullptr);
#endif

  if (!newiobuf) {
    return nullptr;
  }
  Py_CLEAR(input->stringiobuf);
  input->stringiobuf = newiobuf;

  ioobj = IOOOBJECT(newiobuf);
  // This IOBuf is returned to the protocol reader and it doesn't own
  // the underlying buffer. The protocol reader should only call refill
  // when it no longer needs its current IOBuf because the underlying
  // buffer is freed here.
  return folly::IOBuf::wrapBuffer(ioobj->buf + ioobj->pos,
      ioobj->string_size - ioobj->pos);
}

template<typename Reader>
static bool
decodeT(DecodeBuffer *input, PyObject *dec_obj, StructTypeArgs *args,
    int utf8strings) {
  auto refiller = [input] (const uint8_t* remaining,
                           int rlen,
                           int read,
                           int len) {
    return refill(input, remaining, rlen, read, len);
  };
  Reader reader(refiller);
  IOobject *ioobj = IOOOBJECT(input->stringiobuf);
  auto buf = folly::IOBuf::wrapBuffer(ioobj->buf + ioobj->pos,
        ioobj->string_size - ioobj->pos);
  reader.setInput(buf.get());

  try {
    bool ret = decode_struct(&reader, dec_obj, args, utf8strings);
    ioobj = IOOOBJECT(input->stringiobuf);
    ioobj->pos += reader.totalBytesRead();
    return ret;
  } catch (const std::runtime_error& e) {
    PyErr_SetString(PyExc_RuntimeError, e.what());
    return false;
  }
}

static PyObject*
encode(PyObject* /*self*/, PyObject *args, PyObject *kws) {
  PyObject *enc_obj;
  PyObject *spec;
  int utf8strings = 0;
  int protoid = 0;

  static char *kwlist[] = {
    (char*)"enc",
    (char*)"spec",
    (char*)"utf8strings",
    (char*)"protoid",
    nullptr
  };

  if (!PyArg_ParseTupleAndKeywords(args, kws, "OO|ii", kwlist, &enc_obj,
        &spec, &utf8strings, &protoid)) {
    return nullptr;
  }

  if (protoid == 0) {
    return encodeT<BinaryProtocolWriter>(enc_obj, spec, utf8strings);
  } else if (protoid == 2) {
    return encodeT<CompactProtocolWriter>(enc_obj, spec, utf8strings);
  } else {
    PyErr_SetString(PyExc_TypeError, "Unexpected proto id");
    return nullptr;
  }
}

static PyObject*
decode(PyObject* /*self*/, PyObject *args, PyObject *kws) {
  PyObject *dec_obj;
  PyObject *transport;
  PyObject *spec;
  int utf8strings = 0;
  int protoid = 0;
  StructTypeArgs parsedargs;
  DecodeBuffer input = {};

  static char *kwlist[] = {
    (char*)"dec",
    (char*)"transport",
    (char*)"spec",
    (char*)"utf8strings",
    (char*)"protoid",
    nullptr
  };

  if (!PyArg_ParseTupleAndKeywords(args, kws, "OOO|ii", kwlist, &dec_obj,
        &transport, &spec, &utf8strings, &protoid)) {
    return nullptr;
  }

  if (!parse_struct_args(&parsedargs, spec)) {
    return nullptr;
  }

  if (!decode_buffer_from_obj(&input, transport)) {
    return nullptr;
  }

  SCOPE_EXIT {
    free_decodebuf(&input);
  };

  if (protoid == 0) {
    if (!decodeT<BinaryProtocolReaderWithRefill>(&input, dec_obj,
          &parsedargs, utf8strings)) {
      return nullptr;
    }
  } else if (protoid == 2) {
    if (!decodeT<CompactProtocolReaderWithRefill>(&input, dec_obj,
          &parsedargs, utf8strings)) {
      return nullptr;
    }
  } else {
    PyErr_SetString(PyExc_TypeError, "Unexpected proto id");
    return nullptr;
  }

  Py_RETURN_NONE;
}

/* -- PYTHON MODULE SETUP STUFF --- */

static PyMethodDef ThriftFastProtoMethods[] = {

  {"encode", (PyCFunction)encode, METH_VARARGS | METH_KEYWORDS, ""},
  {"decode", (PyCFunction)decode, METH_VARARGS | METH_KEYWORDS, ""},

  {nullptr, nullptr, 0, nullptr}  /* Sentinel */
};

extern "C" {
struct module_state {
  PyObject *error;
};

#if PY_MAJOR_VERSION >= 3
#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))

static int fastproto_traverse(PyObject *m, visitproc visit, void *arg) {
  Py_VISIT(GETSTATE(m)->error);
  return 0;
}

static int fastproto_clear(PyObject *m) {
  Py_CLEAR(GETSTATE(m)->error);
  return 0;
}

static struct PyModuleDef ThriftFastProtoModuleDef = {
  PyModuleDef_HEAD_INIT,
  "thrift.protocol.fastproto",
  nullptr,
  sizeof(struct module_state),
  ThriftFastProtoMethods,
  nullptr,
  fastproto_traverse,
  fastproto_clear,
  nullptr
};

#define INITERROR return nullptr

PyObject*
PyInit_fastproto(void);

PyObject*
PyInit_fastproto(void)
#else
#define INITERROR return
#define GETSTATE(m) (&_state)
static struct module_state _state;

PyMODINIT_FUNC
initfastproto(void);

PyMODINIT_FUNC
initfastproto(void)
#endif
{
  PyObject *module;
  struct module_state *st;
#if PY_MAJOR_VERSION >= 3
  PyObject *iomodule = PyImport_ImportModule("io");
  if (iomodule == nullptr) {
    return nullptr;
  }
  PyObject *bio = PyObject_CallMethod(iomodule, (char*)"BytesIO", (char*)"()");
  if (bio == nullptr) {
    Py_DECREF(iomodule);
    return nullptr;
  }
  BytesIOType = (PyTypeObject*)PyObject_Type(bio);
  module = PyModule_Create(&ThriftFastProtoModuleDef);
#else
  PycString_IMPORT;
  if (PycStringIO == nullptr) return;
  module = Py_InitModule("thrift.protocol.fastproto", ThriftFastProtoMethods);
#endif

  if (module == nullptr) {
    INITERROR;
  }
  st = GETSTATE(module);
  st->error = PyErr_NewException((char*)"fastproto.Error", nullptr, nullptr);
  if (st->error == nullptr) {
    Py_DECREF(module);
    INITERROR;
  }

#if PY_MAJOR_VERSION >= 3
#define INIT_INTERN_STRING(value) \
  do { \
    INTERN_STRING(value) = PyUnicode_InternFromString(#value); \
    if(!INTERN_STRING(value)) return nullptr; \
  } while(0)
#else
#define INIT_INTERN_STRING(value) \
  do { \
    INTERN_STRING(value) = PyString_InternFromString(#value); \
    if(!INTERN_STRING(value)) return; \
  } while(0)
#endif

  INIT_INTERN_STRING(cstringio_buf);
  INIT_INTERN_STRING(cstringio_refill);
#undef INIT_INTERN_STRING

#if PY_MAJOR_VERSION >= 3
  return module;
#endif
}
} // end extern "C"
