/*
 * Copyright 2009-2012 10gen, 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.
 */

/*
 * This file contains C implementations of some of the functions
 * needed by the message module. If possible, these implementations
 * should be used to speed up message creation.
 */

#include "Python.h"

#include "_cbsonmodule.h"
#include "buffer.h"

struct module_state {
    PyObject* _cbson;
};

/* See comments about module initialization in _cbsonmodule.c */
#if PY_MAJOR_VERSION >= 3
#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
#else
#define GETSTATE(m) (&_state)
static struct module_state _state;
#endif

#if PY_MAJOR_VERSION >= 3
#define BYTES_FORMAT_STRING "y#"
#else
#define BYTES_FORMAT_STRING "s#"
#endif

#define DOC_TOO_LARGE_FMT "BSON document too large (%d bytes)" \
                          " - the connected server supports" \
                          " BSON document sizes up to %ld bytes."

/* Get an error class from the pymongo.errors module.
 *
 * Returns a new ref */
static PyObject* _error(char* name) {
    PyObject* error;
    PyObject* errors = PyImport_ImportModule("pymongo.errors");
    if (!errors) {
        return NULL;
    }
    error = PyObject_GetAttrString(errors, name);
    Py_DECREF(errors);
    return error;
}

/* add a lastError message on the end of the buffer.
 * returns 0 on failure */
static int add_last_error(PyObject* self, buffer_t buffer,
                          int request_id, char* ns, int nslen, PyObject* args) {
    struct module_state *state = GETSTATE(self);

    int message_start;
    int document_start;
    int message_length;
    int document_length;
    PyObject* key;
    PyObject* value;
    Py_ssize_t pos = 0;
    PyObject* one;
    char *p = strchr(ns, '.');
    /* Length of the database portion of ns. */
    nslen = p ? (int)(p - ns) : nslen;

    message_start = buffer_save_space(buffer, 4);
    if (message_start == -1) {
        PyErr_NoMemory();
        return 0;
    }
    if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
        !buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"  /* responseTo */
                            "\xd4\x07\x00\x00"  /* opcode */
                            "\x00\x00\x00\x00", /* options */
                            12) ||
        !buffer_write_bytes(buffer,
                            ns, nslen) ||       /* database */
        !buffer_write_bytes(buffer,
                            ".$cmd\x00"         /* collection name */
                            "\x00\x00\x00\x00"  /* skip */
                            "\xFF\xFF\xFF\xFF", /* limit (-1) */
                            14)) {
        return 0;
    }

    /* save space for length */
    document_start = buffer_save_space(buffer, 4);
    if (document_start == -1) {
        PyErr_NoMemory();
        return 0;
    }

    /* getlasterror: 1 */
    if (!(one = PyLong_FromLong(1)))
        return 0;
    if (!write_pair(state->_cbson, buffer, "getlasterror", 12, one, 0, 4, 1)) {
        Py_DECREF(one);
        return 0;
    }
    Py_DECREF(one);

    /* getlasterror options */
    while (PyDict_Next(args, &pos, &key, &value)) {
        if (!decode_and_write_pair(state->_cbson, buffer, key, value, 0, 4, 0)) {
            return 0;
        }
    }

    /* EOD */
    if (!buffer_write_bytes(buffer, "\x00", 1)) {
        return 0;
    }

    message_length = buffer_get_position(buffer) - message_start;
    document_length = buffer_get_position(buffer) - document_start;
    memcpy(buffer_get_buffer(buffer) + message_start, &message_length, 4);
    memcpy(buffer_get_buffer(buffer) + document_start, &document_length, 4);
    return 1;
}

static int init_insert_buffer(buffer_t buffer, int request_id, int options,
                              const char* coll_name, int coll_name_len) {
    /* Save space for message length */
    int length_location = buffer_save_space(buffer, 4);
    if (length_location == -1) {
        PyErr_NoMemory();
        return length_location;
    }
    if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
        !buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"
                            "\xd2\x07\x00\x00",
                            8) ||
        !buffer_write_bytes(buffer, (const char*)&options, 4) ||
        !buffer_write_bytes(buffer,
                            coll_name,
                            coll_name_len + 1)) {
        return -1;
    }
    return length_location;
}

static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) {
    /* Note: As of PyMongo 2.6, this function is no longer used. It
     * is being kept (with tests) for backwards compatibility with 3rd
     * party libraries that may currently be using it, but will likely
     * be removed in a future release. */
    struct module_state *state = GETSTATE(self);

    /* NOTE just using a random number as the request_id */
    int request_id = rand();
    char* collection_name = NULL;
    int collection_name_length;
    PyObject* docs;
    PyObject* doc;
    PyObject* iterator;
    int before, cur_size, max_size = 0;
    int options = 0;
    unsigned char check_keys;
    unsigned char safe;
    unsigned char continue_on_error;
    unsigned char uuid_subtype;
    PyObject* last_error_args;
    buffer_t buffer;
    int length_location, message_length;
    PyObject* result;

    if (!PyArg_ParseTuple(args, "et#ObbObb",
                          "utf-8",
                          &collection_name,
                          &collection_name_length,
                          &docs, &check_keys, &safe,
                          &last_error_args,
                          &continue_on_error, &uuid_subtype)) {
        return NULL;
    }
    if (continue_on_error) {
        options += 1;
    }

    buffer = buffer_new();
    if (!buffer) {
        PyErr_NoMemory();
        PyMem_Free(collection_name);
        return NULL;
    }

    length_location = init_insert_buffer(buffer,
                                         request_id,
                                         options,
                                         collection_name,
                                         collection_name_length);
    if (length_location == -1) {
        PyMem_Free(collection_name);
        buffer_free(buffer);
        return NULL;
    }

    iterator = PyObject_GetIter(docs);
    if (iterator == NULL) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "input is not iterable");
            Py_DECREF(InvalidOperation);
        }
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }
    while ((doc = PyIter_Next(iterator)) != NULL) {
        before = buffer_get_position(buffer);
        if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) {
            Py_DECREF(doc);
            Py_DECREF(iterator);
            buffer_free(buffer);
            PyMem_Free(collection_name);
            return NULL;
        }
        Py_DECREF(doc);
        cur_size = buffer_get_position(buffer) - before;
        max_size = (cur_size > max_size) ? cur_size : max_size;
    }
    Py_DECREF(iterator);

    if (PyErr_Occurred()) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }

    if (!max_size) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert");
            Py_DECREF(InvalidOperation);
        }
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }

    message_length = buffer_get_position(buffer) - length_location;
    memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

    if (safe) {
        if (!add_last_error(self, buffer, request_id, collection_name,
                            collection_name_length, last_error_args)) {
            buffer_free(buffer);
            PyMem_Free(collection_name);
            return NULL;
        }
    }

    PyMem_Free(collection_name);

    /* objectify buffer */
    result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
                           buffer_get_buffer(buffer),
                           buffer_get_position(buffer),
                           max_size);
    buffer_free(buffer);
    return result;
}

PyDoc_STRVAR(_cbson_insert_message_doc,
"Create an insert message to be sent to MongoDB\n\
\n\
Note: As of PyMongo 2.6, this function is no longer used. It\n\
is being kept (with tests) for backwards compatibility with 3rd\n\
party libraries that may currently be using it, but will likely\n\
be removed in a future release.");

static PyObject* _cbson_update_message(PyObject* self, PyObject* args) {
    /* NOTE just using a random number as the request_id */
    struct module_state *state = GETSTATE(self);

    int request_id = rand();
    char* collection_name = NULL;
    int collection_name_length;
    int before, cur_size, max_size = 0;
    PyObject* doc;
    PyObject* spec;
    unsigned char multi;
    unsigned char upsert;
    unsigned char safe;
    unsigned char check_keys;
    unsigned char uuid_subtype;
    PyObject* last_error_args;
    int options;
    buffer_t buffer;
    int length_location, message_length;
    PyObject* result;

    if (!PyArg_ParseTuple(args, "et#bbOObObb",
                          "utf-8",
                          &collection_name,
                          &collection_name_length,
                          &upsert, &multi, &spec, &doc, &safe,
                          &last_error_args, &check_keys, &uuid_subtype)) {
        return NULL;
    }

    options = 0;
    if (upsert) {
        options += 1;
    }
    if (multi) {
        options += 2;
    }
    buffer = buffer_new();
    if (!buffer) {
        PyErr_NoMemory();
        PyMem_Free(collection_name);
        return NULL;
    }

    // save space for message length
    length_location = buffer_save_space(buffer, 4);
    if (length_location == -1) {
        PyMem_Free(collection_name);
        PyErr_NoMemory();
        return NULL;
    }
    if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
        !buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"
                            "\xd1\x07\x00\x00"
                            "\x00\x00\x00\x00",
                            12) ||
        !buffer_write_bytes(buffer,
                            collection_name,
                            collection_name_length + 1) ||
        !buffer_write_bytes(buffer, (const char*)&options, 4)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }

    before = buffer_get_position(buffer);
    if (!write_dict(state->_cbson, buffer, spec, 0, uuid_subtype, 1)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }
    max_size = buffer_get_position(buffer) - before;

    before = buffer_get_position(buffer);
    if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }
    cur_size = buffer_get_position(buffer) - before;
    max_size = (cur_size > max_size) ? cur_size : max_size;

    message_length = buffer_get_position(buffer) - length_location;
    memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

    if (safe) {
        if (!add_last_error(self, buffer, request_id, collection_name,
                            collection_name_length, last_error_args)) {
            buffer_free(buffer);
            PyMem_Free(collection_name);
            return NULL;
        }
    }

    PyMem_Free(collection_name);

    /* objectify buffer */
    result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
                           buffer_get_buffer(buffer),
                           buffer_get_position(buffer),
                           max_size);
    buffer_free(buffer);
    return result;
}

static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
    /* NOTE just using a random number as the request_id */
    struct module_state *state = GETSTATE(self);

    int request_id = rand();
    unsigned int options;
    char* collection_name = NULL;
    int collection_name_length;
    int begin, cur_size, max_size = 0;
    int num_to_skip;
    int num_to_return;
    PyObject* query;
    PyObject* field_selector = Py_None;
    unsigned char uuid_subtype = 3;
    buffer_t buffer;
    int length_location, message_length;
    PyObject* result;

    if (!PyArg_ParseTuple(args, "Iet#iiO|Ob",
                          &options,
                          "utf-8",
                          &collection_name,
                          &collection_name_length,
                          &num_to_skip, &num_to_return,
                          &query, &field_selector, &uuid_subtype)) {
        return NULL;
    }
    buffer = buffer_new();
    if (!buffer) {
        PyErr_NoMemory();
        PyMem_Free(collection_name);
        return NULL;
    }

    // save space for message length
    length_location = buffer_save_space(buffer, 4);
    if (length_location == -1) {
        PyMem_Free(collection_name);
        PyErr_NoMemory();
        return NULL;
    }
    if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
        !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) ||
        !buffer_write_bytes(buffer, (const char*)&options, 4) ||
        !buffer_write_bytes(buffer, collection_name,
                            collection_name_length + 1) ||
        !buffer_write_bytes(buffer, (const char*)&num_to_skip, 4) ||
        !buffer_write_bytes(buffer, (const char*)&num_to_return, 4)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }

    begin = buffer_get_position(buffer);
    if (!write_dict(state->_cbson, buffer, query, 0, uuid_subtype, 1)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }
    max_size = buffer_get_position(buffer) - begin;

    if (field_selector != Py_None) {
        begin = buffer_get_position(buffer);
        if (!write_dict(state->_cbson, buffer, field_selector, 0, uuid_subtype, 1)) {
            buffer_free(buffer);
            PyMem_Free(collection_name);
            return NULL;
        }
        cur_size = buffer_get_position(buffer) - begin;
        max_size = (cur_size > max_size) ? cur_size : max_size;
    }

    PyMem_Free(collection_name);

    message_length = buffer_get_position(buffer) - length_location;
    memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

    /* objectify buffer */
    result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
                           buffer_get_buffer(buffer),
                           buffer_get_position(buffer),
                           max_size);
    buffer_free(buffer);
    return result;
}

static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) {
    /* NOTE just using a random number as the request_id */
    int request_id = rand();
    char* collection_name = NULL;
    int collection_name_length;
    int num_to_return;
    long long cursor_id;
    buffer_t buffer;
    int length_location, message_length;
    PyObject* result;

    if (!PyArg_ParseTuple(args, "et#iL",
                          "utf-8",
                          &collection_name,
                          &collection_name_length,
                          &num_to_return,
                          &cursor_id)) {
        return NULL;
    }
    buffer = buffer_new();
    if (!buffer) {
        PyErr_NoMemory();
        PyMem_Free(collection_name);
        return NULL;
    }

    // save space for message length
    length_location = buffer_save_space(buffer, 4);
    if (length_location == -1) {
        PyMem_Free(collection_name);
        PyErr_NoMemory();
        return NULL;
    }
    if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
        !buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"
                            "\xd5\x07\x00\x00"
                            "\x00\x00\x00\x00", 12) ||
        !buffer_write_bytes(buffer,
                            collection_name,
                            collection_name_length + 1) ||
        !buffer_write_bytes(buffer, (const char*)&num_to_return, 4) ||
        !buffer_write_bytes(buffer, (const char*)&cursor_id, 8)) {
        buffer_free(buffer);
        PyMem_Free(collection_name);
        return NULL;
    }

    PyMem_Free(collection_name);

    message_length = buffer_get_position(buffer) - length_location;
    memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

    /* objectify buffer */
    result = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id,
                           buffer_get_buffer(buffer),
                           buffer_get_position(buffer));
    buffer_free(buffer);
    return result;
}

static void
_set_document_too_large(int size, long max) {
    PyObject* InvalidDocument = _error("InvalidDocument");
    if (InvalidDocument) {
#if PY_MAJOR_VERSION >= 3
        PyObject* error = PyUnicode_FromFormat(DOC_TOO_LARGE_FMT, size, max);
#else
        PyObject* error = PyString_FromFormat(DOC_TOO_LARGE_FMT, size, max);
#endif
        if (error) {
            PyErr_SetObject(InvalidDocument, error);
            Py_DECREF(error);
        }
        Py_DECREF(InvalidDocument);
    }
}

static PyObject* _cbson_do_batched_insert(PyObject* self, PyObject* args) {
    struct module_state *state = GETSTATE(self);

    /* NOTE just using a random number as the request_id */
    int request_id = rand();
    int options = 0;
    int length_location, message_length;
    int collection_name_length;
    char* collection_name = NULL;
    PyObject* docs;
    PyObject* doc;
    PyObject* iterator;
    PyObject* client;
    PyObject* last_error_args;
    PyObject* result;
    PyObject* max_bson_size_obj;
    PyObject* max_message_size_obj;
    PyObject* send_message_result;
    unsigned char check_keys;
    unsigned char safe;
    unsigned char continue_on_error;
    unsigned char uuid_subtype;
    unsigned char empty = 1;
    long max_bson_size;
    long max_message_size;
    buffer_t buffer;
    PyObject *exc_type = NULL, *exc_value = NULL, *exc_trace = NULL;

    if (!PyArg_ParseTuple(args, "et#ObbObbO",
                          "utf-8",
                          &collection_name,
                          &collection_name_length,
                          &docs, &check_keys, &safe,
                          &last_error_args,
                          &continue_on_error,
                          &uuid_subtype, &client)) {
        return NULL;
    }
    if (continue_on_error) {
        options += 1;
    }

    max_bson_size_obj = PyObject_GetAttrString(client, "max_bson_size");
#if PY_MAJOR_VERSION >= 3
    max_bson_size = PyLong_AsLong(max_bson_size_obj);
#else
    max_bson_size = PyInt_AsLong(max_bson_size_obj);
#endif
    Py_XDECREF(max_bson_size_obj);
    if (max_bson_size == -1) {
        PyMem_Free(collection_name);
        return NULL;
    }

    max_message_size_obj = PyObject_GetAttrString(client, "max_message_size");
#if PY_MAJOR_VERSION >= 3
    max_message_size = PyLong_AsLong(max_message_size_obj);
#else
    max_message_size = PyInt_AsLong(max_message_size_obj);
#endif
    Py_XDECREF(max_message_size_obj);
    if (max_message_size == -1) {
        PyMem_Free(collection_name);
        return NULL;
    }

    buffer = buffer_new();
    if (!buffer) {
        PyErr_NoMemory();
        PyMem_Free(collection_name);
        return NULL;
    }

    length_location = init_insert_buffer(buffer,
                                         request_id,
                                         options,
                                         collection_name,
                                         collection_name_length);
    if (length_location == -1) {
        goto insertfail;
    }

    iterator = PyObject_GetIter(docs);
    if (iterator == NULL) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "input is not iterable");
            Py_DECREF(InvalidOperation);
        }
        goto insertfail;
    }
    while ((doc = PyIter_Next(iterator)) != NULL) {
        int before = buffer_get_position(buffer);
        int cur_size;
        empty = 0;
        if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) {
            Py_DECREF(doc);
            goto iterfail;
        }
        Py_DECREF(doc);

        cur_size = buffer_get_position(buffer) - before;
        if (cur_size > max_bson_size) {
            _set_document_too_large(cur_size, max_bson_size);
            goto iterfail;
        }

        /* We have enough data, send this batch. */
        if (buffer_get_position(buffer) > max_message_size) {
            int new_request_id = rand();
            int message_start;
            PyObject* send_gle = Py_False;
            buffer_t new_buffer = buffer_new();
            if (!new_buffer) {
                PyErr_NoMemory();
                goto iterfail;
            }
            message_start = init_insert_buffer(new_buffer,
                                               new_request_id,
                                               options,
                                               collection_name,
                                               collection_name_length);
            if (message_start == -1) {
                buffer_free(new_buffer);
                goto iterfail;
            }

            /* Copy the overflow encoded document into the new buffer. */
            if (!buffer_write_bytes(new_buffer,
                (const char*)buffer_get_buffer(buffer) + before, cur_size)) {
                buffer_free(new_buffer);
                goto iterfail;
            }

            /* Roll back to the beginning of this document. */
            buffer_update_position(buffer, before);
            message_length = buffer_get_position(buffer) - length_location;
            memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

            /* If we are doing unacknowledged writes *and* continue_on_error
             * is True it's pointless (and slower) to send GLE. */
            if (safe || !continue_on_error) {
                send_gle = Py_True;
                if (!add_last_error(self, buffer, request_id, collection_name,
                                    collection_name_length, last_error_args)) {
                    buffer_free(new_buffer);
                    goto iterfail;
                }
            }
            /* Objectify buffer */
            result = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id,
                                   buffer_get_buffer(buffer),
                                   buffer_get_position(buffer));
            buffer_free(buffer);
            buffer = new_buffer;
            request_id = new_request_id;
            length_location = message_start;

            send_message_result = PyObject_CallMethod(client, "_send_message",
                                                      "NO", result, send_gle);

            if (!send_message_result) {
                PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
                PyObject* OperationFailure;
                PyErr_Fetch(&etype, &evalue, &etrace);
                OperationFailure = _error("OperationFailure");
                if (OperationFailure) {
                    if (PyErr_GivenExceptionMatches(etype, OperationFailure)) {
                        if (!safe || continue_on_error) {
                            Py_DECREF(OperationFailure);
                            if (!safe) {
                                /* We're doing unacknowledged writes and
                                 * continue_on_error is False. Just return. */
                                Py_DECREF(etype);
                                Py_XDECREF(evalue);
                                Py_XDECREF(etrace);
                                Py_DECREF(iterator);
                                buffer_free(buffer);
                                PyMem_Free(collection_name);
                                Py_RETURN_NONE;
                            }
                            /* continue_on_error is True, store the error
                             * details to re-raise after the final batch */
                            Py_XDECREF(exc_type);
                            Py_XDECREF(exc_value);
                            Py_XDECREF(exc_trace);
                            exc_type = etype;
                            exc_value = evalue;
                            exc_trace = etrace;
                            continue;
                        }
                    }
                    Py_DECREF(OperationFailure);
                }
                /* This isn't OperationFailure, we couldn't
                 * import OperationFailure, or we are doing
                 * acknowledged writes. Re-raise immediately. */
                PyErr_Restore(etype, evalue, etrace);
                goto iterfail;
            } else {
                Py_DECREF(send_message_result);
            }
        }
    }
    Py_DECREF(iterator);

    if (PyErr_Occurred()) {
        goto insertfail;
    }

    if (empty) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert");
            Py_DECREF(InvalidOperation);
        }
        goto insertfail;
    }

    message_length = buffer_get_position(buffer) - length_location;
    memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);

    if (safe) {
        if (!add_last_error(self, buffer, request_id, collection_name,
                            collection_name_length, last_error_args)) {
            goto insertfail;
        }
    }

    PyMem_Free(collection_name);

    /* objectify buffer */
    result = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id,
                           buffer_get_buffer(buffer),
                           buffer_get_position(buffer));
    buffer_free(buffer);

    /* Send the last (or only) batch */
    send_message_result = PyObject_CallMethod(client, "_send_message", "NN",
                                              result,
                                              PyBool_FromLong((long)safe));

    if (!send_message_result) {
        Py_XDECREF(exc_type);
        Py_XDECREF(exc_value);
        Py_XDECREF(exc_trace);
        return NULL;
    } else {
        Py_DECREF(send_message_result);
    }

    if (exc_type) {
        /* Re-raise any previously stored exception
         * due to continue_on_error being True */
        PyErr_Restore(exc_type, exc_value, exc_trace);
        return NULL;
    }

    Py_RETURN_NONE;

iterfail:
    Py_DECREF(iterator);
insertfail:
    Py_XDECREF(exc_type);
    Py_XDECREF(exc_value);
    Py_XDECREF(exc_trace);
    buffer_free(buffer);
    PyMem_Free(collection_name);
    return NULL;
}

static PyObject*
_send_write_command(PyObject* client, buffer_t buffer,
                    int lst_len_loc, int cmd_len_loc, unsigned char* errors) {

    PyObject* msg;
    PyObject* result;
    PyObject* ok;

    int request_id = rand();
    int position = buffer_get_position(buffer);
    int length = position - lst_len_loc - 1;
    memcpy(buffer_get_buffer(buffer) + lst_len_loc, &length, 4);
    length = position - cmd_len_loc;
    memcpy(buffer_get_buffer(buffer) + cmd_len_loc, &length, 4);
    memcpy(buffer_get_buffer(buffer), &position, 4);
    memcpy(buffer_get_buffer(buffer) + 4, &request_id, 4);

    /* objectify buffer */
    msg = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id,
                        buffer_get_buffer(buffer),
                        buffer_get_position(buffer));
    if (!msg)
        return NULL;

    /* Send the current batch */
    result = PyObject_CallMethod(client, "_send_message",
                                 "NOO", msg, Py_True, Py_True);
    if (result && PyDict_GetItemString(result, "writeErrors"))
        *errors = 1;
    return result;
}

static buffer_t
_command_buffer_new(char* ns, int ns_len) {
    buffer_t buffer;
    if (!(buffer = buffer_new())) {
        PyErr_NoMemory();
        return NULL;
    }
    /* Save space for message length and request id */
    if ((buffer_save_space(buffer, 8)) == -1) {
        PyErr_NoMemory();
        buffer_free(buffer);
        return NULL;
    }
    if (!buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"  /* responseTo */
                            "\xd4\x07\x00\x00"  /* opcode */
                            "\x00\x00\x00\x00", /* options */
                            12) ||
        !buffer_write_bytes(buffer,
                            ns, ns_len + 1) ||  /* namespace */
        !buffer_write_bytes(buffer,
                            "\x00\x00\x00\x00"  /* skip */
                            "\xFF\xFF\xFF\xFF", /* limit (-1) */
                            8)) {
        buffer_free(buffer);
        return NULL;
    }
    return buffer;
}

#define _INSERT 0
#define _UPDATE 1
#define _DELETE 2

static PyObject*
_cbson_do_batched_write_command(PyObject* self, PyObject* args) {
    struct module_state *state = GETSTATE(self);

    long max_bson_size;
    long max_cmd_size;
    long idx_offset = 0;
    int idx = 0;
    int cmd_len_loc;
    int lst_len_loc;
    int ns_len;
    int ordered;
    char *ns = NULL;
    PyObject* max_bson_size_obj;
    PyObject* command;
    PyObject* doc;
    PyObject* docs;
    PyObject* client;
    PyObject* iterator;
    PyObject* result;
    PyObject* results;
    unsigned char op;
    unsigned char check_keys;
    unsigned char uuid_subtype;
    unsigned char empty = 1;
    unsigned char errors = 0;
    buffer_t buffer;

    if (!PyArg_ParseTuple(args, "et#bOObbO", "utf-8",
                          &ns, &ns_len, &op, &command, &docs,
                          &check_keys, &uuid_subtype, &client)) {
        return NULL;
    }

    max_bson_size_obj = PyObject_GetAttrString(client, "max_bson_size");
#if PY_MAJOR_VERSION >= 3
    max_bson_size = PyLong_AsLong(max_bson_size_obj);
#else
    max_bson_size = PyInt_AsLong(max_bson_size_obj);
#endif
    Py_XDECREF(max_bson_size_obj);
    if (max_bson_size == -1) {
        PyMem_Free(ns);
        return NULL;
    }
    /*
     * Max BSON object size + 16k - 2 bytes for ending NUL bytes
     * XXX: This should come from the server - SERVER-10643
     */
    max_cmd_size = max_bson_size + 16382;

    /* Default to True */
    ordered = !((PyDict_GetItemString(command, "ordered")) == Py_False);

    if (!(results = PyList_New(0))) {
        PyMem_Free(ns);
        return NULL;
    }

    if (!(buffer = _command_buffer_new(ns, ns_len))) {
        PyMem_Free(ns);
        Py_DECREF(results);
        return NULL;
    }

    PyMem_Free(ns);

    /* Position of command document length */
    cmd_len_loc = buffer_get_position(buffer);
    if (!write_dict(state->_cbson, buffer, command, 0, uuid_subtype, 0)) {
        goto cmdfail;
    }

    /* Write type byte for array */
    *(buffer_get_buffer(buffer) + (buffer_get_position(buffer) - 1)) = 0x4;

    switch (op) {
    case _INSERT:
        {
            if (!buffer_write_bytes(buffer, "documents\x00", 10))
                goto cmdfail;
            break;
        }
    case _UPDATE:
        {
            /* MongoDB does key validation for update. */
            check_keys = 0;
            if (!buffer_write_bytes(buffer, "updates\x00", 8))
                goto cmdfail;
            break;
        }
    case _DELETE:
        {
            /* Never check keys in a delete command. */
            check_keys = 0;
            if (!buffer_write_bytes(buffer, "deletes\x00", 8))
                goto cmdfail;
            break;
        }
    default:
        {
            PyObject* InvalidOperation = _error("InvalidOperation");
            if (InvalidOperation) {
                PyErr_SetString(InvalidOperation, "Unknown command");
                Py_DECREF(InvalidOperation);
            }
            goto cmdfail;
        }
    }

    /* Save space for list document */
    lst_len_loc = buffer_save_space(buffer, 4);
    if (lst_len_loc == -1) {
        PyErr_NoMemory();
        goto cmdfail;
    }

    iterator = PyObject_GetIter(docs);
    if (iterator == NULL) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "input is not iterable");
            Py_DECREF(InvalidOperation);
        }
        goto cmdfail;
    }
    while ((doc = PyIter_Next(iterator)) != NULL) {
        int sub_doc_begin = buffer_get_position(buffer);
        int cur_doc_begin;
        int cur_size;
        char key[16];
        empty = 0;
        INT2STRING(key, idx);
        if (!buffer_write_bytes(buffer, "\x03", 1) ||
            !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) {
            Py_DECREF(doc);
            goto cmditerfail;
        }
        cur_doc_begin = buffer_get_position(buffer);
        if (!write_dict(state->_cbson, buffer, doc,
                        check_keys, uuid_subtype, 1)) {
            Py_DECREF(doc);
            goto cmditerfail;
        }
        Py_DECREF(doc);

        /* We have enough data, maybe send this batch. */
        if (buffer_get_position(buffer) > max_cmd_size) {
            buffer_t new_buffer;
            cur_size = buffer_get_position(buffer) - cur_doc_begin;

            /* This single document is too large for the command. */
            if (!idx) {
                if (op == _INSERT) {
                    _set_document_too_large(cur_size, max_bson_size);
                } else {
                    PyObject* InvalidDocument = _error("InvalidDocument");
                    if (InvalidDocument) {
                        /*
                         * There's nothing intelligent we can say
                         * about size for update and remove.
                         */
                        PyErr_SetString(InvalidDocument,
                                        "command document too large");
                        Py_DECREF(InvalidDocument);
                    }
                }
                goto cmditerfail;
            }

            if (!(new_buffer = buffer_new())) {
                PyErr_NoMemory();
                goto cmditerfail;
            }
            /* New buffer including the current overflow document */
            if (!buffer_write_bytes(new_buffer,
                (const char*)buffer_get_buffer(buffer), lst_len_loc + 5) ||
                !buffer_write_bytes(new_buffer, "0\x00", 2) ||
                !buffer_write_bytes(new_buffer,
                (const char*)buffer_get_buffer(buffer) + cur_doc_begin, cur_size)) {
                buffer_free(new_buffer);
                goto cmditerfail;
            }
            /*
             * Roll the existing buffer back to the beginning
             * of the last document encoded.
             */
            buffer_update_position(buffer, sub_doc_begin);

            if (!buffer_write_bytes(buffer, "\x00\x00", 2))
                goto cmditerfail;

            result = _send_write_command(client, buffer,
                                         lst_len_loc, cmd_len_loc, &errors);

            buffer_free(buffer);
            buffer = new_buffer;

            if (!result)
                goto cmditerfail;

#if PY_MAJOR_VERSION >= 3
            result = Py_BuildValue("NN",
                                   PyLong_FromLong(idx_offset), result);
#else
            result = Py_BuildValue("NN",
                                   PyInt_FromLong(idx_offset), result);
#endif
            if (!result)
                goto cmditerfail;

            PyList_Append(results, result);
            Py_DECREF(result);

            if (errors && ordered) {
                Py_DECREF(iterator);
                buffer_free(buffer);
                return results;
            }
            idx_offset += idx;
            idx = 0;
        }
        idx += 1;
    }
    Py_DECREF(iterator);

    if (PyErr_Occurred()) {
        goto cmdfail;
    }

    if (empty) {
        PyObject* InvalidOperation = _error("InvalidOperation");
        if (InvalidOperation) {
            PyErr_SetString(InvalidOperation, "cannot do an empty bulk write");
            Py_DECREF(InvalidOperation);
        }
        goto cmdfail;
    }

    if (!buffer_write_bytes(buffer, "\x00\x00", 2))
        goto cmdfail;

    result = _send_write_command(client, buffer,
                                 lst_len_loc, cmd_len_loc, &errors);
    if (!result)
        goto cmdfail;

#if PY_MAJOR_VERSION >= 3
    result = Py_BuildValue("NN", PyLong_FromLong(idx_offset), result);
#else
    result = Py_BuildValue("NN", PyInt_FromLong(idx_offset), result);
#endif
    if (!result)
        goto cmdfail;

    buffer_free(buffer);

    PyList_Append(results, result);
    Py_DECREF(result);
    return results;

cmditerfail:
    Py_DECREF(iterator);
cmdfail:
    Py_DECREF(results);
    buffer_free(buffer);
    return NULL;
}

static PyMethodDef _CMessageMethods[] = {
    {"_insert_message", _cbson_insert_message, METH_VARARGS,
     _cbson_insert_message_doc},
    {"_update_message", _cbson_update_message, METH_VARARGS,
     "create an update message to be sent to MongoDB"},
    {"_query_message", _cbson_query_message, METH_VARARGS,
     "create a query message to be sent to MongoDB"},
    {"_get_more_message", _cbson_get_more_message, METH_VARARGS,
     "create a get more message to be sent to MongoDB"},
    {"_do_batched_insert", _cbson_do_batched_insert, METH_VARARGS,
     "insert a batch of documents, splitting the batch as needed"},
    {"_do_batched_write_command", _cbson_do_batched_write_command, METH_VARARGS,
     "execute a batch of insert, update, or delete commands"},
    {NULL, NULL, 0, NULL}
};

#if PY_MAJOR_VERSION >= 3
#define INITERROR return NULL
static int _cmessage_traverse(PyObject *m, visitproc visit, void *arg) {
    Py_VISIT(GETSTATE(m)->_cbson);
    return 0;
}

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

static struct PyModuleDef moduledef = {
        PyModuleDef_HEAD_INIT,
        "_cmessage",
        NULL,
        sizeof(struct module_state),
        _CMessageMethods,
        NULL,
        _cmessage_traverse,
        _cmessage_clear,
        NULL
};

PyMODINIT_FUNC
PyInit__cmessage(void)
#else
#define INITERROR return
PyMODINIT_FUNC
init_cmessage(void)
#endif
{
    PyObject *_cbson;
    PyObject *c_api_object;
    PyObject *m;
    struct module_state *state;

    /* Store a reference to the _cbson module since it's needed to call some
     * of its functions
     */
    _cbson = PyImport_ImportModule("bson._cbson");
    if (_cbson == NULL) {
        INITERROR;
    }

    /* Import C API of _cbson
     * The header file accesses _cbson_API to call the functions
     */
    c_api_object = PyObject_GetAttrString(_cbson, "_C_API");
    if (c_api_object == NULL) {
        Py_DECREF(_cbson);
        INITERROR;
    }
#if PY_VERSION_HEX >= 0x03010000
    _cbson_API = (void **)PyCapsule_GetPointer(c_api_object, "_cbson._C_API");
#else
    _cbson_API = (void **)PyCObject_AsVoidPtr(c_api_object);
#endif
    if (_cbson_API == NULL) {
        Py_DECREF(c_api_object);
        Py_DECREF(_cbson);
        INITERROR;
    }

#if PY_MAJOR_VERSION >= 3
    m = PyModule_Create(&moduledef);
#else
    m = Py_InitModule("_cmessage", _CMessageMethods);
#endif
    if (m == NULL) {
        Py_DECREF(c_api_object);
        Py_DECREF(_cbson);
        INITERROR;
    }

    state = GETSTATE(m);
    state->_cbson = _cbson;

    Py_DECREF(c_api_object);

#if PY_MAJOR_VERSION >= 3
    return m;
#endif
}