#include "tree.h"

#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <errno.h>
#include <inttypes.h>
#include <netdb.h>
#include <stdbool.h>
#include <stdio.h>

#define LOCAL

#ifdef __GNUC__
#  define UNUSED(x) UNUSED_ ## x __attribute__((__unused__))
#else
#  define UNUSED(x) UNUSED_ ## x
#endif

#define NETWORK_BIT_VALUE(network, current_bit)                    \
    (network)->bytes[((network)->max_depth0 - (current_bit)) >> 3] \
    & (1U << (~((network)->max_depth0 - (current_bit)) & 7))

/* This is also defined in MaxMind::DB::Common but we don't want to have to
 * fetch it every time we need it. */
#define DATA_SECTION_SEPARATOR_SIZE (16)

#define SHA1_KEY_LENGTH (27)

#define NETWORK_IS_IPV6(network) (127 == network->max_depth0)

typedef struct freeze_args_s {
    int fd;
    char *filename;
    HV *data_hash;
} freeze_args_s;

typedef struct thawed_network_s {
    MMDBW_network_s *network;
    MMDBW_record_s *record;
} thawed_network_s;

typedef struct encode_args_s {
    PerlIO *output_io;
    SV *root_data_type;
    SV *serializer;
    HV *data_pointer_cache;
} encode_args_s;

/* *INDENT-OFF* */
/* --prototypes automatically generated by dev-bin/regen-prototypes.pl - don't remove this comment */
LOCAL void insert_resolved_network(MMDBW_tree_s *tree, MMDBW_network_s *network,
                                   SV *key_sv, SV *data);
LOCAL const char *store_data_in_tree(MMDBW_tree_s *tree,
                                     const char *const key,
                                     SV *data_sv);
LOCAL const char *increment_data_reference_count(MMDBW_tree_s *tree,
                                                 const char *const key);
LOCAL void set_stored_data_in_tree(MMDBW_tree_s *tree,
                                   const char *const key,
                                   SV *data_sv);
LOCAL void decrement_data_reference_count(MMDBW_tree_s *tree,
                                          const char *const key);
LOCAL MMDBW_network_s resolve_network(MMDBW_tree_s *tree,
                                      const char *const ipstr,
                                      const uint8_t prefix_length);
LOCAL void free_network(MMDBW_network_s *network);
LOCAL void insert_record_for_network(MMDBW_tree_s *tree,
                                     MMDBW_network_s *network,
                                     MMDBW_record_s *new_record,
                                     bool merge_record_collisions);
LOCAL bool merge_records(MMDBW_tree_s *tree,
                         MMDBW_network_s *network,
                         MMDBW_record_s *new_record,
                         MMDBW_record_s *record_to_set);
LOCAL void merge_hash(HV *from, HV *to);
LOCAL MMDBW_node_s *find_node_for_network(MMDBW_tree_s *tree,
                                          MMDBW_network_s *network,
                                          uint8_t *current_bit,
                                          MMDBW_node_s *(if_not_node)(
                                              MMDBW_tree_s *tree,
                                              MMDBW_record_s *record));
LOCAL MMDBW_node_s *return_null(
    MMDBW_tree_s *UNUSED(tree), MMDBW_record_s *UNUSED(record));
LOCAL MMDBW_node_s *new_node_from_record(MMDBW_tree_s *tree,
                                         MMDBW_record_s *record);
LOCAL void free_node_and_subnodes(MMDBW_tree_s *tree, MMDBW_node_s *node);
LOCAL void free_record_value(MMDBW_tree_s *tree, MMDBW_record_s *record);
LOCAL void assign_node_numbers(MMDBW_tree_s *tree);
LOCAL void assign_node_number(MMDBW_tree_s *tree, MMDBW_node_s *node,
                              uint128_t UNUSED(network),
                              uint8_t UNUSED(depth), void *UNUSED(args));
LOCAL void freeze_node(MMDBW_tree_s *tree, MMDBW_node_s *node,
                       uint128_t network, uint8_t depth, void *void_args);
LOCAL void freeze_data_record(MMDBW_tree_s *UNUSED(tree),
                              uint128_t network, uint8_t depth,
                              const char *key,
                              freeze_args_s *args);
LOCAL void freeze_to_fd(freeze_args_s *args, void *data, size_t size);
LOCAL void freeze_data_to_fd(int fd, MMDBW_tree_s *tree);
LOCAL SV *freeze_hash(HV *hash);
LOCAL uint8_t thaw_uint8(uint8_t **buffer);
LOCAL uint32_t thaw_uint32(uint8_t **buffer);
LOCAL thawed_network_s *thaw_network(MMDBW_tree_s *tree, uint8_t **buffer);
LOCAL uint8_t *thaw_bytes(uint8_t **buffer, size_t size);
LOCAL uint128_t thaw_uint128(uint8_t **buffer);
LOCAL STRLEN thaw_strlen(uint8_t **buffer);
LOCAL const char *thaw_data_key(uint8_t **buffer);
LOCAL HV *thaw_data_hash(SV *data_to_decode);
LOCAL void encode_node(MMDBW_tree_s *tree, MMDBW_node_s *node,
                       uint128_t UNUSED(network),
                       uint8_t UNUSED(depth), void *void_args);
LOCAL void check_record_sanity(MMDBW_node_s *node, MMDBW_record_s *record,
                               char *side);
LOCAL uint32_t record_value_as_number(MMDBW_tree_s *tree,
                                      MMDBW_record_s *record,
                                      encode_args_s * args);
LOCAL void iterate_tree(MMDBW_tree_s *tree,
                        MMDBW_node_s *node,
                        uint128_t network,
                        const uint8_t depth,
                        bool depth_first,
                        void *args,
                        void(callback) (MMDBW_tree_s *tree,
                                        MMDBW_node_s *node,
                                        uint128_t network,
                                        const uint8_t depth,
                                        void *args));
LOCAL SV *key_for_data(SV * data);
LOCAL void dwarn(SV *thing);
LOCAL void *checked_malloc(size_t size);
LOCAL void checked_write(int fd, char *filename, void *buffer,
                         ssize_t count);
LOCAL void checked_perlio_read(PerlIO * io, void *buffer,
                               SSize_t size);
LOCAL void check_perlio_result(SSize_t result, SSize_t expected,
                               char *op);
/* --prototypes end - don't remove this comment-- */
/* *INDENT-ON* */

MMDBW_tree_s *new_tree(const uint8_t ip_version, uint8_t record_size,
                       bool merge_record_collisions)
{
    MMDBW_tree_s *tree = checked_malloc(sizeof(MMDBW_tree_s));

    /* XXX - check for 4 or 6 */
    tree->ip_version = ip_version;
    /* XXX - check for 24, 28, or 32 */
    tree->record_size = record_size;
    tree->merge_record_collisions = merge_record_collisions;
    tree->data_table = NULL;
    tree->is_finalized = false;
    tree->is_aliased = false;
    tree->root_node = new_node(tree);
    tree->node_count = 0;

    return tree;
}

int insert_network(MMDBW_tree_s *tree, const char *const ipstr,
                   const uint8_t prefix_length, SV *key, SV *data)
{
    MMDBW_network_s network = resolve_network(tree, ipstr, prefix_length);

    if (tree->ip_version == 4 && NETWORK_IS_IPV6((&network))) {
        free_network(&network);
        return -1;
    }

    insert_resolved_network(tree, &network, key, data);

    free_network(&network);

    return 0;
}

LOCAL void insert_resolved_network(MMDBW_tree_s *tree, MMDBW_network_s *network,
                                   SV *key_sv, SV *data)
{
    const char *const key =
        store_data_in_tree(tree, SvPVbyte_nolen(key_sv), data);
    MMDBW_record_s new_record = {
        .type    = MMDBW_RECORD_TYPE_DATA,
        .value   = {
            .key = key
        }
    };

    insert_record_for_network(tree, network, &new_record,
                              tree->merge_record_collisions);
}

LOCAL const char *store_data_in_tree(MMDBW_tree_s *tree,
                                     const char *const key,
                                     SV *data_sv)
{
    const char *const new_key = increment_data_reference_count(tree, key);
    set_stored_data_in_tree(tree, key, data_sv);

    return new_key;
}

LOCAL const char *increment_data_reference_count(MMDBW_tree_s *tree,
                                                 const char *const key)
{
    MMDBW_data_hash_s *data = NULL;
    HASH_FIND(hh, tree->data_table, key, SHA1_KEY_LENGTH, data);

    /* We allow this possibility as we need to create the record separately
       from updating the data when thawing */
    if (NULL == data) {
        data = checked_malloc(sizeof(MMDBW_data_hash_s));
        data->reference_count = 0;

        data->data_sv = NULL;

        data->key = checked_malloc(SHA1_KEY_LENGTH + 1);
        strcpy((char *)data->key, key);

        HASH_ADD_KEYPTR(hh, tree->data_table, data->key, SHA1_KEY_LENGTH, data);
    }
    data->reference_count++;

    return data->key;
}


LOCAL void set_stored_data_in_tree(MMDBW_tree_s *tree,
                                   const char *const key,
                                   SV *data_sv)
{
    MMDBW_data_hash_s *data = NULL;
    HASH_FIND(hh, tree->data_table, key, SHA1_KEY_LENGTH, data);

    if (NULL == data) {
        croak("Attempt to set unknown data record in tree");
    }

    if (NULL != data->data_sv) {
        return;
    }

    SvREFCNT_inc_simple_void_NN(data_sv);
    data->data_sv = data_sv;
}

LOCAL void decrement_data_reference_count(MMDBW_tree_s *tree,
                                          const char *const key)
{
    MMDBW_data_hash_s *data = NULL;
    HASH_FIND(hh, tree->data_table, key, SHA1_KEY_LENGTH, data);

    if (NULL == data) {
        croak("Attempt to remove data that does not exist from tree");
    }

    data->reference_count--;
    if (0 == data->reference_count) {
        HASH_DEL(tree->data_table, data);
        SvREFCNT_dec(data->data_sv);
        free((char *)data->key);
        free(data);
    }
}


/* XXX - this is mostly copied from libmaxminddb - can we somehow share this code? */
LOCAL MMDBW_network_s resolve_network(MMDBW_tree_s *tree,
                                      const char *const ipstr,
                                      const uint8_t prefix_length)
{
    struct addrinfo ai_hints;
    ai_hints.ai_socktype = 0;
    ai_hints.ai_protocol = 0;
    if (tree->ip_version == 6 || NULL != strchr(ipstr, ':')) {
        ai_hints.ai_flags = AI_NUMERICHOST | AI_V4MAPPED;
        ai_hints.ai_family = AF_INET6;
    } else {
        ai_hints.ai_flags = AI_NUMERICHOST;
        ai_hints.ai_family = AF_INET;
    }

    struct addrinfo *addresses;
    int status = getaddrinfo(ipstr, NULL, &ai_hints, &addresses);
    if (status) {
        croak("Bad IP address: %s - %s\n", ipstr, gai_strerror(status));
    }

    int family = addresses->ai_addr->sa_family;
    uint8_t *bytes;
    if (family == AF_INET) {
        bytes = checked_malloc(4);
        memcpy(bytes,
               &((struct sockaddr_in *)addresses->ai_addr)->sin_addr.s_addr,
               4);
    } else {
        bytes = checked_malloc(16);
        memcpy(
            bytes,
            ((struct sockaddr_in6 *)addresses->ai_addr)->sin6_addr.s6_addr,
            16);
    }

    MMDBW_network_s network = {
        .bytes          = bytes,
        .prefix_length  = prefix_length,
        .max_depth0     = (family == AF_INET ? 31 : 127),
        .address_string = ipstr,
    };

    freeaddrinfo(addresses);

    return network;
}

LOCAL void free_network(MMDBW_network_s *network)
{
    free((char *)network->bytes);
}

struct network {
    const char *const ipstr;
    const uint8_t prefix_length;
};

static struct network ipv4_aliases[] = {
    {
        .ipstr = "::ffff:0:0",
        .prefix_length = 96
    },
    {
        .ipstr = "2001::",
        .prefix_length = 32
    },
    {
        .ipstr = "2002::",
        .prefix_length = 16
    }
};

void alias_ipv4_networks(MMDBW_tree_s *tree)
{
    if (tree->ip_version == 4) {
        return;
    }
    if (tree->is_aliased) {
        return;
    }

    MMDBW_network_s ipv4_root_network = resolve_network(tree, "::0.0.0.0", 96);

    uint8_t current_bit;
    MMDBW_node_s *ipv4_root_node_parent =
        find_node_for_network(tree, &ipv4_root_network, &current_bit,
                              &return_null);
    /* If the current_bit is not 32 then we found some node further up the
     * tree that would eventually lead to that network. This means that there
     * are no IPv4 addresses in the tree, so there's nothing to alias. */
    if (32 != current_bit) {
        free_network(&ipv4_root_network);
        return;
    }

    if (MMDBW_RECORD_TYPE_NODE != ipv4_root_node_parent->left_record.type) {
        free_network(&ipv4_root_network);
        return;
    }

    MMDBW_node_s *ipv4_root_node =
        ipv4_root_node_parent->left_record.value.node;
    for (int i = 0; i <= 2; i++) {
        MMDBW_network_s alias_network =
            resolve_network(tree, ipv4_aliases[i].ipstr,
                            ipv4_aliases[i].prefix_length);

        MMDBW_node_s *last_node_for_alias = find_node_for_network(
            tree, &alias_network, &current_bit, &new_node_from_record);
        if (NETWORK_BIT_VALUE(&alias_network, current_bit)) {
            free_record_value(tree, &(last_node_for_alias->right_record));
            last_node_for_alias->right_record.type = MMDBW_RECORD_TYPE_ALIAS;
            last_node_for_alias->right_record.value.node = ipv4_root_node;
        } else {
            free_record_value(tree, &(last_node_for_alias->left_record));
            last_node_for_alias->left_record.type = MMDBW_RECORD_TYPE_ALIAS;
            last_node_for_alias->left_record.value.node = ipv4_root_node;
        }

        free_network(&alias_network);
    }

    free_network(&ipv4_root_network);
}

LOCAL void insert_record_for_network(MMDBW_tree_s *tree,
                                     MMDBW_network_s *network,
                                     MMDBW_record_s *new_record,
                                     bool merge_record_collisions)
{
    uint8_t current_bit;
    MMDBW_node_s *node_to_set =
        find_node_for_network(tree, network, &current_bit,
                              &new_node_from_record);

    MMDBW_record_s *record_to_set, *other_record;
    if (NETWORK_BIT_VALUE(network, current_bit)) {
        record_to_set = &(node_to_set->right_record);
        other_record = &(node_to_set->left_record);
    } else {
        record_to_set = &(node_to_set->left_record);
        other_record = &(node_to_set->right_record);
    }

    if (merge_record_collisions &&
        MMDBW_RECORD_TYPE_DATA == new_record->type) {

        if (merge_records(tree, network, new_record, record_to_set)) {
            return;
        }
    }

    /* If this record we're about to insert is a data record, and the other
     * record in the node also has the same data, then we instead want to
     * insert a single data record in this node's parent. We do this by
     * inserting the new record for the parent network, which we can calculate
     * quite easily by subtracting 1 from this network's prefix length. */
    if (MMDBW_RECORD_TYPE_DATA == new_record->type
        && MMDBW_RECORD_TYPE_DATA == other_record->type
        ) {

        const char *const new_key = new_record->value.key;
        const char *const other_key = other_record->value.key;

        if (strlen(new_key) == strlen(other_key)
            && 0 == strcmp(new_key, other_key)) {

            size_t bytes_length = NETWORK_IS_IPV6(network) ? 16 : 4;
            uint8_t *bytes = checked_malloc(bytes_length);
            memcpy(bytes, network->bytes, bytes_length);

            uint8_t parent_prefix_length = network->prefix_length - 1;
            MMDBW_network_s parent_network = {
                .bytes          = bytes,
                .prefix_length  = parent_prefix_length,
                .max_depth0     = network->max_depth0,
                .address_string = network->address_string,
            };

            /* We don't need to merge record collisions in this insert as
             * we have already merged the new record with the existing
             * record
             */
            insert_record_for_network(tree, &parent_network, new_record, false);
            free_network(&parent_network);
            return;
        }
    }

    free_record_value(tree, record_to_set);

    record_to_set->type = new_record->type;
    if (MMDBW_RECORD_TYPE_DATA == new_record->type) {
        record_to_set->value.key = new_record->value.key;
    } else if (MMDBW_RECORD_TYPE_NODE == new_record->type ||
               MMDBW_RECORD_TYPE_ALIAS == new_record->type) {
        record_to_set->value.node = new_record->value.node;
    }

    return;
}

LOCAL bool merge_records(MMDBW_tree_s *tree,
                         MMDBW_network_s *network,
                         MMDBW_record_s *new_record,
                         MMDBW_record_s *record_to_set)
{
    if (MMDBW_RECORD_TYPE_NODE == record_to_set->type ||
        MMDBW_RECORD_TYPE_ALIAS == record_to_set->type) {

        if (network->prefix_length > network->max_depth0) {
            croak("Something is very wrong. Prefix length is too long.");
        }

        /* We increment the count as we are turning one record into two */
        increment_data_reference_count(tree, new_record->value.key);

        uint8_t new_prefix_length = network->prefix_length + 1;

        MMDBW_network_s left = {
            .bytes          = network->bytes,
            .prefix_length  = new_prefix_length,
            .max_depth0     = network->max_depth0,
            .address_string = network->address_string,
        };

        MMDBW_record_s new_left_record = {
            .type    = new_record->type,
            .value   = {
                .key = new_record->value.key
            }
        };

        insert_record_for_network(tree, &left, &new_left_record, true);

        size_t bytes_length = NETWORK_IS_IPV6(network) ? 16 : 4;
        uint8_t right_bytes[bytes_length];
        memcpy(&right_bytes, network->bytes, bytes_length);

        right_bytes[ (new_prefix_length - 1) / 8]
            |= 1 << ((network->max_depth0 + 1 - new_prefix_length) % 8);

        char right_address_string[NETWORK_IS_IPV6(network) ? INET6_ADDRSTRLEN :
                                  INET_ADDRSTRLEN];
        inet_ntop(NETWORK_IS_IPV6(network) ? AF_INET6 : AF_INET,
                  &right_bytes,
                  right_address_string,
                  sizeof(right_address_string));

        MMDBW_network_s right = {
            .bytes          = (const uint8_t *const)&right_bytes,
            .prefix_length  = new_prefix_length,
            .max_depth0     = network->max_depth0,
            .address_string = right_address_string,
        };

        MMDBW_record_s new_right_record = {
            .type    = new_record->type,
            .value   = {
                .key = new_record->value.key
            }
        };

        insert_record_for_network(tree, &right, &new_right_record, true);

        /* There's no need continuing with the original record as the relevant
         * data has already been inserted further down the tree by the code
         * above. */
        return true;
    }
    /* This must come before the node pruning code in
       insert_record_for_network, as we only want to prune nodes where the
       merged record matches. */
    else if (MMDBW_RECORD_TYPE_DATA == record_to_set->type) {
        SV *merged = merge_hashes_for_keys(tree,
                                           record_to_set->value.key,
                                           new_record->value.key,
                                           network);

        SV *key_sv = key_for_data(merged);
        const char *const new_key =
            store_data_in_tree(tree, SvPVbyte_nolen(key_sv), merged);
        SvREFCNT_dec(key_sv);

        /* The ref count was incremented in store_data_in_tree */
        SvREFCNT_dec(merged);

        decrement_data_reference_count(tree, new_record->value.key);
        new_record->value.key = new_key;
    }

    return false;
}

SV *merge_hashes_for_keys(MMDBW_tree_s *tree, const char *const key_from,
                          const char *const key_into, MMDBW_network_s *network)
{
    SV *data_from = data_for_key(tree, key_from);
    SV *data_into = data_for_key(tree, key_into);

    if (!(SvROK(data_from) && SvROK(data_into)
          && SvTYPE(SvRV(data_from)) == SVt_PVHV
          && SvTYPE(SvRV(data_into)) == SVt_PVHV)) {
        /* We added key_into earlier during insert_resolved_network, so we got to
           make sure here that it's removed again after we decide to not actually
           store this network. It might be nicer to not insert anything into the
           tree until we're sure we really want to. */
        decrement_data_reference_count(tree, key_into);
        croak(
            "Cannot merge data records unless both records are hashes - inserting %s/%"
            PRIu8,
            network->address_string, network->prefix_length);
    }

    HV *hash_from = (HV *)SvRV(data_from);
    HV *hash_into = (HV *)SvRV(data_into);
    HV *hash_new = newHV();

    merge_hash(hash_into, hash_new);
    merge_hash(hash_from, hash_new);

    return newRV_noinc((SV *)hash_new);
}

LOCAL void merge_hash(HV *from, HV *to)
{
    (void)hv_iterinit(from);
    HE *he;
    while (NULL != (he = hv_iternext(from))) {
        STRLEN key_length;
        const char *const key = HePV(he, key_length);
        U32 hash = HeHASH(he);
        if (hv_exists(to, key, key_length)) {
            continue;
        }

        SV *value = HeVAL(he);
        SvREFCNT_inc_simple_void_NN(value);
        (void)hv_store(to, key, key_length, value, hash);
    }

    return;
}

SV *lookup_ip_address(MMDBW_tree_s *tree, const char *const ipstr)
{
    MMDBW_network_s network =
        resolve_network(tree, ipstr, tree->ip_version == 6 ? 128 : 32);

    uint8_t current_bit;
    MMDBW_node_s *node_for_address =
        find_node_for_network(tree, &network, &current_bit, &return_null);

    MMDBW_record_s record_for_address;
    if (NETWORK_BIT_VALUE(&network, current_bit)) {
        record_for_address = node_for_address->right_record;
    } else {
        record_for_address = node_for_address->left_record;
    }

    free_network(&network);

    if (MMDBW_RECORD_TYPE_NODE == record_for_address.type ||
        MMDBW_RECORD_TYPE_ALIAS == record_for_address.type) {
        croak(
            "WTF - found a node or alias record for an address lookup - %s - current_bit = %"
            PRIu8,
            ipstr, current_bit);
        return &PL_sv_undef;
    } else if (MMDBW_RECORD_TYPE_EMPTY == record_for_address.type) {
        return &PL_sv_undef;
    } else {
        return newSVsv(data_for_key(tree, record_for_address.value.key));
    }
}

LOCAL MMDBW_node_s *find_node_for_network(MMDBW_tree_s *tree,
                                          MMDBW_network_s *network,
                                          uint8_t *current_bit,
                                          MMDBW_node_s *(if_not_node)(
                                              MMDBW_tree_s *tree,
                                              MMDBW_record_s *record))
{
    MMDBW_node_s *node = tree->root_node;
    uint8_t last_bit = network->max_depth0 - (network->prefix_length - 1);

    for (*current_bit = network->max_depth0; *current_bit > last_bit;
         (*current_bit)--) {

        int next_is_right = NETWORK_BIT_VALUE(network, *current_bit);
        MMDBW_record_s *record =
            next_is_right
            ? &(node->right_record)
            : &(node->left_record);

        MMDBW_node_s *next_node;
        if (MMDBW_RECORD_TYPE_NODE == record->type ||
            MMDBW_RECORD_TYPE_ALIAS == record->type) {
            next_node = record->value.node;
        } else {
            next_node = if_not_node(tree, record);
            if (NULL == next_node) {
                return node;
            }

            record->type = MMDBW_RECORD_TYPE_NODE;
            record->value.node = next_node;
        }

        node = next_node;
    }

    return node;
}

LOCAL MMDBW_node_s *return_null(
    MMDBW_tree_s *UNUSED(tree), MMDBW_record_s *UNUSED(record))
{
    return NULL;
}

LOCAL MMDBW_node_s *new_node_from_record(MMDBW_tree_s *tree,
                                         MMDBW_record_s *record)
{
    MMDBW_node_s *node = new_node(tree);
    if (MMDBW_RECORD_TYPE_DATA == record->type) {
        /* We only need to increment the reference count once as we are
           replacing the parent record */
        increment_data_reference_count(tree, record->value.key);

        node->left_record.type = MMDBW_RECORD_TYPE_DATA;
        node->left_record.value.key = record->value.key;

        node->right_record.type = MMDBW_RECORD_TYPE_DATA;
        node->right_record.value.key = record->value.key;
    }

    return node;
}

MMDBW_node_s *new_node(MMDBW_tree_s *tree)
{
    MMDBW_node_s *node = checked_malloc(sizeof(MMDBW_node_s));

    node->number = 0;
    node->left_record.type = node->right_record.type = MMDBW_RECORD_TYPE_EMPTY;

    tree->is_finalized = false;

    return node;
}

LOCAL void free_node_and_subnodes(MMDBW_tree_s *tree, MMDBW_node_s *node)
{
    free_record_value(tree, &(node->left_record));
    free_record_value(tree, &(node->right_record));

    free(node);
}

LOCAL void free_record_value(MMDBW_tree_s *tree, MMDBW_record_s *record)
{
    if (MMDBW_RECORD_TYPE_NODE == record->type) {
        free_node_and_subnodes(tree, record->value.node);
    }

    if (MMDBW_RECORD_TYPE_DATA == record->type) {
        decrement_data_reference_count(tree, record->value.key);
    }

    /* We do not follow MMDBW_RECORD_TYPE_ALIAS nodes */
}


void finalize_tree(MMDBW_tree_s *tree)
{
    if (tree->is_finalized) {
        return;
    }

    assign_node_numbers(tree);
    tree->is_finalized = true;
}

LOCAL void assign_node_numbers(MMDBW_tree_s *tree)
{
    tree->node_count = 0;
    start_iteration(tree, false, (void *)NULL, &assign_node_number);
}

LOCAL void assign_node_number(MMDBW_tree_s *tree, MMDBW_node_s *node,
                              uint128_t UNUSED(network),
                              uint8_t UNUSED(depth), void *UNUSED(args))
{
    node->number = tree->node_count++;
    return;
}

/* 16 bytes for an IP address, 1 byte for the prefix length */
#define FROZEN_RECORD_MAX_SIZE (16 + 1 + SHA1_KEY_LENGTH)
#define FROZEN_NODE_MAX_SIZE (FROZEN_RECORD_MAX_SIZE * 2)

/* 17 bytes of NULLs followed by something that cannot be an SHA1 key are a
   clear indicator that there are no more frozen networks in the buffer. */
#define SEVENTEEN_NULLS "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"
#define FREEZE_SEPARATOR "not an SHA1 key"
/* We subtract 1 as we treat this as a sequence of bytes rather than a null terminated
   string. */
#define FREEZE_SEPARATOR_LENGTH (sizeof(FREEZE_SEPARATOR) - 1)

void freeze_tree(MMDBW_tree_s *tree, char *filename, char *frozen_params,
                 size_t frozen_params_size)
{
    finalize_tree(tree);

    int fd = open(filename, O_CREAT | O_TRUNC | O_RDWR, (mode_t)0644);
    if (-1 == fd) {
        croak("Could not open file %s: %s", filename, strerror(errno));
    }

    freeze_args_s args = {
        .fd       = fd,
        .filename = filename,
    };

    freeze_to_fd(&args, &frozen_params_size, 4);
    freeze_to_fd(&args, frozen_params, frozen_params_size);

    start_iteration(tree, false, (void *)&args, &freeze_node);

    freeze_to_fd(&args, SEVENTEEN_NULLS, 17);
    freeze_to_fd(&args, FREEZE_SEPARATOR, FREEZE_SEPARATOR_LENGTH);

    freeze_data_to_fd(fd, tree);

    if (-1 == close(fd)) {
        croak("Could not close file %s: %s", filename, strerror(errno));
    }

    /* When the hash is _freed_, Perl decrements the ref count for each value
     * so we don't need to mess with them. */
    SvREFCNT_dec((SV *)args.data_hash);
}

LOCAL void freeze_node(MMDBW_tree_s *tree, MMDBW_node_s *node,
                       uint128_t network, uint8_t depth, void *void_args)
{
    freeze_args_s *args = (freeze_args_s *)void_args;

    const uint8_t max_depth0 = tree->ip_version == 6 ? 127 : 31;
    const uint8_t next_depth = depth + 1;

    if (MMDBW_RECORD_TYPE_DATA == node->left_record.type) {
        freeze_data_record(tree, network, next_depth,
                           node->left_record.value.key, args);
    }

    if (MMDBW_RECORD_TYPE_DATA == node->right_record.type) {
        uint128_t right_network =
            FLIP_NETWORK_BIT(network, max_depth0, depth);
        freeze_data_record(tree, right_network, next_depth,
                           node->right_record.value.key, args);
    }
}

LOCAL void freeze_data_record(MMDBW_tree_s *UNUSED(tree),
                              uint128_t network, uint8_t depth,
                              const char *key,
                              freeze_args_s *args)
{
    /* It'd save some space to shrink this to 4 bytes for IPv4-only trees, but
     * that would also complicated thawing quite a bit. */
    freeze_to_fd(args, &network, 16);
    freeze_to_fd(args, &(depth), 1);
    freeze_to_fd(args, (char *)key, SHA1_KEY_LENGTH);
}

LOCAL void freeze_to_fd(freeze_args_s *args, void *data, size_t size)
{
    checked_write(args->fd, args->filename, data, size);
}

LOCAL void freeze_data_to_fd(int fd, MMDBW_tree_s *tree)
{
    HV *data_hash = newHV();

    MMDBW_data_hash_s *item, *tmp;
    HASH_ITER(hh, tree->data_table, item, tmp) {
        SvREFCNT_inc_simple_void_NN(item->data_sv);
        (void)hv_store(data_hash, item->key, SHA1_KEY_LENGTH, item->data_sv, 0);
    }

    SV *frozen_data = freeze_hash(data_hash);
    STRLEN frozen_data_size;
    char *frozen_data_chars = SvPV(frozen_data, frozen_data_size);

    ssize_t written = write(fd, &frozen_data_size, sizeof(STRLEN));
    if (-1 == written) {
        croak("Could not write frozen data size to file: %s", strerror(errno));
    }
    if (written != sizeof(STRLEN)) {
        croak("Could not write frozen data size to file: %zd != %zu", written,
              sizeof(STRLEN));
    }

    written = write(fd, frozen_data_chars, frozen_data_size);
    if (-1 == written) {
        croak("Could not write frozen data size to file: %s", strerror(errno));
    }
    if (written != (ssize_t)frozen_data_size) {
        croak("Could not write frozen data to file: %zd != %zu", written,
              frozen_data_size);
    }

    SvREFCNT_dec(frozen_data);
    SvREFCNT_dec((SV *)data_hash);
}

LOCAL SV *freeze_hash(HV *hash)
{
    dSP;
    ENTER;
    SAVETMPS;

    SV *hashref = sv_2mortal(newRV_inc((SV *)hash));

    PUSHMARK(SP);
    EXTEND(SP, 1);
    PUSHs(hashref);
    PUTBACK;

    int count = call_pv("Sereal::Encoder::encode_sereal", G_SCALAR);

    SPAGAIN;

    if (count != 1) {
        croak("Expected 1 item back from Sereal::Encoder::encode_sereal call");
    }

    SV *frozen = POPs;
    if (!SvPOK(frozen)) {
        croak(
            "The Sereal::Encoder::encode_sereal sub returned an SV which is not SvPOK!");
    }

    /* The SV will be mortal so it's about to lose a ref with the FREETMPS
       call below. */
    SvREFCNT_inc_simple_void_NN(frozen);

    PUTBACK;
    FREETMPS;
    LEAVE;

    return frozen;
}

MMDBW_tree_s *thaw_tree(char *filename, uint32_t initial_offset,
                        uint8_t ip_version, uint8_t record_size,
                        bool merge_record_collisions)
{
    int fd = open(filename, O_RDONLY, 0);
    if (-1 == fd) {
        croak("Could not open file %s: %s", filename, strerror(errno));
    }

    struct stat fileinfo;
    if (-1 == fstat(fd, &fileinfo)) {
        close(fd);
        croak("Could not stat file: %s: %s", filename, strerror(errno));
    }

    uint8_t *buffer =
        (uint8_t *)mmap(NULL, fileinfo.st_size, PROT_READ, MAP_SHARED, fd,
                        0);
    close(fd);

    buffer += initial_offset;

    MMDBW_tree_s *tree = new_tree(ip_version, record_size,
                                  merge_record_collisions);

    thawed_network_s *thawed;
    while (NULL != (thawed = thaw_network(tree, &buffer))) {
        if (MMDBW_RECORD_TYPE_DATA == thawed->record->type) {
            const char *key = increment_data_reference_count(
                tree, thawed->record->value.key);

            /* insert_record_for_network reuses the key. We want it to use
               the same copy as used in the data hash. */
            free((char *)thawed->record->value.key);
            thawed->record->value.key = key;
        }
        insert_record_for_network(tree, thawed->network, thawed->record,
                                  tree->merge_record_collisions);
        free_network(thawed->network);
        free(thawed->network);
        free(thawed->record);
        free(thawed);
    }

    STRLEN frozen_data_size = thaw_strlen(&buffer);

    /* per perlapi newSVpvn copies the string */
    SV *data_to_decode =
        sv_2mortal(newSVpvn((char *)buffer, frozen_data_size));
    HV *data_hash = thaw_data_hash(data_to_decode);

    hv_iterinit(data_hash);
    char *key;
    I32 keylen;
    SV *value;
    while (NULL != (value = hv_iternextsv(data_hash, &key, &keylen))) {
        set_stored_data_in_tree(tree, key, value);
    }

    SvREFCNT_dec((SV *)data_hash);

    finalize_tree(tree);

    return tree;
}

LOCAL uint8_t thaw_uint8(uint8_t **buffer)
{
    uint8_t value;
    memcpy(&value, *buffer, 1);
    *buffer += 1;
    return value;
}

LOCAL uint32_t thaw_uint32(uint8_t **buffer)
{
    uint32_t value;
    memcpy(&value, *buffer, 4);
    *buffer += 4;
    return value;
}

LOCAL thawed_network_s *thaw_network(MMDBW_tree_s *tree, uint8_t **buffer)
{
    uint128_t start_ip = thaw_uint128(buffer);
    uint8_t prefix_length = thaw_uint8(buffer);

    if (0 == start_ip && 0 == prefix_length) {
        uint8_t *maybe_separator = thaw_bytes(buffer, FREEZE_SEPARATOR_LENGTH);
        if (memcmp(maybe_separator, FREEZE_SEPARATOR,
                   FREEZE_SEPARATOR_LENGTH) == 0) {

            free(maybe_separator);
            return NULL;
        }

        croak("Found a ::0/0 network but that should never happen!");
    }

    uint8_t *start_ip_bytes = (uint8_t *)&start_ip;
    uint8_t temp;
    for (int i = 0; i < 8; i++) {
        temp = start_ip_bytes[i];
        start_ip_bytes[i] = start_ip_bytes[15 - i];
        start_ip_bytes[15 - i] = temp;
    }

    thawed_network_s *thawed = checked_malloc(sizeof(thawed_network_s));

    uint8_t *bytes;
    if (tree->ip_version == 4) {
        bytes = checked_malloc(4);
        memcpy(bytes, start_ip_bytes + 12, 4);
    } else {
        bytes = checked_malloc(16);
        memcpy(bytes, &start_ip, 16);
    }

    MMDBW_network_s network = {
        .bytes          = bytes,
        .prefix_length  = prefix_length,
        .max_depth0     = 4 == tree->ip_version ? 31 : 127,
        .address_string = "thawed network",
    };

    thawed->network = checked_malloc(sizeof(MMDBW_network_s));
    memcpy(thawed->network, &network, sizeof(MMDBW_network_s));

    MMDBW_record_s *record = checked_malloc(sizeof(MMDBW_record_s));
    record->type = MMDBW_RECORD_TYPE_DATA;

    record->value.key = thaw_data_key(buffer);

    thawed->record = record;

    return thawed;
}

LOCAL uint8_t *thaw_bytes(uint8_t **buffer, size_t size)
{
    uint8_t *value = checked_malloc(size);
    memcpy(value, *buffer, size);
    *buffer += size;
    return value;
}

LOCAL uint128_t thaw_uint128(uint8_t **buffer)
{
    uint128_t value;
    memcpy(&value, *buffer, 16);
    *buffer += 16;
    return value;
}

LOCAL STRLEN thaw_strlen(uint8_t **buffer)
{
    STRLEN value;
    memcpy(&value, *buffer, sizeof(STRLEN));
    *buffer += sizeof(STRLEN);
    return value;
}

LOCAL const char *thaw_data_key(uint8_t **buffer)
{
    char *value = checked_malloc(SHA1_KEY_LENGTH + 1);
    memcpy(value, *buffer, SHA1_KEY_LENGTH);
    *buffer += SHA1_KEY_LENGTH;
    value[SHA1_KEY_LENGTH] = '\0';
    return (const char *)value;
}

LOCAL HV *thaw_data_hash(SV *data_to_decode)
{
    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);
    EXTEND(SP, 1);
    PUSHs(data_to_decode);
    PUTBACK;

    int count = call_pv("Sereal::Decoder::decode_sereal", G_SCALAR);

    SPAGAIN;

    if (count != 1) {
        croak("Expected 1 item back from Sereal::Decoder::decode_sereal call");
    }

    SV *thawed = POPs;
    if (!SvROK(thawed)) {
        croak(
            "The Sereal::Decoder::decode_sereal sub returned an SV which is not SvROK!");
    }

    SV *data_hash = SvREFCNT_inc_simple_NN(SvRV(thawed));

    PUTBACK;
    FREETMPS;
    LEAVE;

    return (HV *)data_hash;
}

void write_search_tree(MMDBW_tree_s *tree, SV *output,
                       const bool alias_ipv6,
                       SV *root_data_type, SV *serializer)
{
    if (alias_ipv6) {
        alias_ipv4_networks(tree);
    }

    finalize_tree(tree);

    /* This is a gross way to get around the fact that with C function
     * pointers we can't easily pass different params to different
     * callbacks. */
    encode_args_s args = {
        .output_io          = IoOFP(sv_2io(output)),
        .root_data_type     = root_data_type,
        .serializer         = serializer,
        .data_pointer_cache = newHV()
    };

    start_iteration(tree, false, (void *)&args, &encode_node);

    /* When the hash is _freed_, Perl decrements the ref count for each value
     * so we don't need to mess with them. */
    SvREFCNT_dec((SV *)args.data_pointer_cache);

    return;
}

LOCAL void encode_node(MMDBW_tree_s *tree, MMDBW_node_s *node,
                       uint128_t UNUSED(network),
                       uint8_t UNUSED(depth), void *void_args)
{
    encode_args_s *args = (encode_args_s *)void_args;

    check_record_sanity(node, &(node->left_record), "left");
    check_record_sanity(node, &(node->right_record), "right");

    uint32_t left =
        htonl(record_value_as_number(tree, &(node->left_record), args));
    uint32_t right =
        htonl(record_value_as_number(tree, &(node->right_record), args));

    uint8_t *left_bytes = (uint8_t *)&left;
    uint8_t *right_bytes = (uint8_t *)&right;

    if (24 == tree->record_size) {
        check_perlio_result(
            PerlIO_printf(args->output_io, "%c%c%c%c%c%c",
                          left_bytes[1], left_bytes[2], left_bytes[3],
                          right_bytes[1], right_bytes[2],
                          right_bytes[3]),
            6, "PerlIO_printf");
    } else if (28 == tree->record_size) {
        check_perlio_result(
            PerlIO_printf(args->output_io, "%c%c%c%c%c%c%c",
                          left_bytes[1], left_bytes[2],
                          left_bytes[3],
                          (left_bytes[0] <<
                           4) | (right_bytes[0] & 15),
                          right_bytes[1], right_bytes[2],
                          right_bytes[3]),
            7, "PerlIO_printf");
    } else {
        check_perlio_result(
            PerlIO_printf(args->output_io, "%c%c%c%c%c%c%c%c",
                          left_bytes[0], left_bytes[1],
                          left_bytes[2], left_bytes[3],
                          right_bytes[0], right_bytes[1],
                          right_bytes[2], right_bytes[3]),
            8, "PerlIO_printf");
    }
}

/* Note that for data records, we will ensure that the key they contain does
 * match a data record in the record_value_as_number() subroutine. */
LOCAL void check_record_sanity(MMDBW_node_s *node, MMDBW_record_s *record,
                               char *side)
{
    if (MMDBW_RECORD_TYPE_NODE == record->type) {
        if (record->value.node->number == node->number) {
            croak("%s record of node %" PRIu32 " points to the same node",
                  side, node->number);
        }

        if (record->value.node->number < node->number) {
            croak(
                "%s record of node %" PRIu32 " points to a node  number (%"
                PRIu32
                ")",
                side, node->number, record->value.node->number);
        }
    }

    if (MMDBW_RECORD_TYPE_ALIAS == record->type) {
        if (0 == record->value.node->number) {
            croak("%s record of node %" PRIu32 " is an alias to node 0",
                  side, node->number);
        }
    }
}

LOCAL uint32_t record_value_as_number(MMDBW_tree_s *tree,
                                      MMDBW_record_s *record,
                                      encode_args_s * args)
{
    uint32_t record_value;

    if (MMDBW_RECORD_TYPE_EMPTY == record->type) {
        record_value = tree->node_count;
    } else if (MMDBW_RECORD_TYPE_NODE == record->type ||
               MMDBW_RECORD_TYPE_ALIAS == record->type) {
        record_value = record->value.node->number;
    } else {
        SV **cache_record =
            hv_fetch(args->data_pointer_cache, record->value.key,
                     SHA1_KEY_LENGTH, 0);
        if (cache_record) {
            /* It is ok to return this without the size check below as it
               would have already croaked when it was inserted if it was too
               big. */
            return SvIV(*cache_record);
        }

        SV *data = newSVsv(data_for_key(tree, record->value.key));
        if (!SvOK(data)) {
            croak("No data associated with key - %s", record->value.key);
        }

        dSP;
        ENTER;
        SAVETMPS;

        PUSHMARK(SP);
        EXTEND(SP, 5);
        PUSHs(args->serializer);
        PUSHs(args->root_data_type);
        mPUSHs(data);
        PUSHs(&PL_sv_undef);
        mPUSHp(record->value.key, strlen(record->value.key));
        PUTBACK;

        int count = call_method("store_data", G_SCALAR);

        SPAGAIN;

        if (count != 1) {
            croak("Expected 1 item back from ->store_data() call");
        }

        SV *rval = POPs;
        if (!(SvIOK(rval) || SvUOK(rval))) {
            croak(
                "The serializer's store_data() method returned an SV which is not SvIOK or SvUOK!");
        }
        uint32_t position = (uint32_t )SvUV(rval);

        PUTBACK;
        FREETMPS;
        LEAVE;

        record_value = position + tree->node_count +
                       DATA_SECTION_SEPARATOR_SIZE;

        SV *value = newSViv(record_value);
        (void)hv_store(args->data_pointer_cache, record->value.key,
                       SHA1_KEY_LENGTH, value, 0);
    }

    if (record_value > MAX_RECORD_VALUE(tree->record_size)) {
        croak(
            "Node value of %" PRIu32 " exceeds the record size of %" PRIu8
            " bits",
            record_value, tree->record_size);
    }

    return record_value;
}

void start_iteration(MMDBW_tree_s *tree,
                     bool depth_first,
                     void *args,
                     void(callback) (MMDBW_tree_s *tree,
                                     MMDBW_node_s *node,
                                     uint128_t network,
                                     uint8_t depth,
                                     void *args))
{
    uint128_t network = 0;
    uint8_t depth = 0;

    iterate_tree(tree, tree->root_node, network, depth, depth_first,
                 args, callback);

    return;
}

LOCAL void iterate_tree(MMDBW_tree_s *tree,
                        MMDBW_node_s *node,
                        uint128_t network,
                        const uint8_t depth,
                        bool depth_first,
                        void *args,
                        void(callback) (MMDBW_tree_s *tree,
                                        MMDBW_node_s *node,
                                        uint128_t network,
                                        const uint8_t depth,
                                        void *args))
{
    if (depth > 127) {
        croak(
            "Depth during iteration is greater than 127! The tree is wonky.\n");
    }

    if (!depth_first) {
        callback(tree, node, network, depth, args);
    }

    if (MMDBW_RECORD_TYPE_NODE == node->left_record.type) {
        iterate_tree(tree,
                     node->left_record.value.node,
                     network,
                     depth + 1,
                     depth_first,
                     args,
                     callback);
    }

    if (depth_first) {
        callback(tree, node, network, depth, args);
    }

    const uint8_t max_depth0 = tree->ip_version == 6 ? 127 : 31;
    if (MMDBW_RECORD_TYPE_NODE == node->right_record.type) {
        iterate_tree(tree,
                     node->right_record.value.node,
                     FLIP_NETWORK_BIT(network, max_depth0, depth),
                     depth + 1,
                     depth_first,
                     args,
                     callback);
    }
}

LOCAL SV *key_for_data(SV * data)
{
    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);
    EXTEND(SP, 1);
    PUSHs(data);
    PUTBACK;

    const char *const sub = "MaxMind::DB::Writer::Util::key_for_data";
    int count = call_pv(sub, G_SCALAR);

    SPAGAIN;

    if (count != 1) {
        croak("Expected 1 item back from %s() call", sub);
    }

    SV *key = POPs;
    SvREFCNT_inc_simple_void_NN(key);

    PUTBACK;
    FREETMPS;
    LEAVE;

    return key;
}

SV *data_for_key(MMDBW_tree_s *tree, const char *const key)
{
    MMDBW_data_hash_s *data = NULL;
    HASH_FIND(hh, tree->data_table, key, strlen(key), data);

    if (NULL != data) {
        return data->data_sv;
    } else {
        return &PL_sv_undef;
    }
}

void free_tree(MMDBW_tree_s *tree)
{
    finalize_tree(tree);

    free_node_and_subnodes(tree, tree->root_node);

    int hash_count = HASH_COUNT(tree->data_table);
    if (0 != hash_count) {
        croak("%d elements left in data table after freeing all nodes!",
              hash_count);
    }

    free(tree);
}

const char *record_type_name(int record_type)
{
    return MMDBW_RECORD_TYPE_EMPTY == record_type
           ? "empty"
           : MMDBW_RECORD_TYPE_NODE == record_type
           ? "node"
           : MMDBW_RECORD_TYPE_ALIAS == record_type
           ? "alias"
           : "data";
}

static SV *module;
LOCAL void dwarn(SV *thing)
{
    if (NULL == module) {
        module = newSVpv("Devel::Dwarn", 0);
        load_module(PERL_LOADMOD_NOIMPORT, module, NULL);
    }

    dSP;
    ENTER;
    SAVETMPS;

    PUSHMARK(SP);
    EXTEND(SP, 1);
    PUSHs(thing);
    PUTBACK;

    (void)call_pv("Devel::Dwarn::Dwarn", G_VOID);

    SPAGAIN;

    PUTBACK;
    FREETMPS;
    LEAVE;
}

void warn_hex(uint8_t digest[16], char *where)
{
    char *hex = md5_as_hex(digest);
    fprintf(stderr, "MD5 = %s (%s)\n", hex, where);
    free(hex);
}

char *md5_as_hex(uint8_t digest[16])
{
    char *hex = checked_malloc(33);
    for (int i = 0; i < 16; ++i) {
        sprintf(&hex[i * 2], "%02x", digest[i]);
    }

    return hex;
}

LOCAL void *checked_malloc(size_t size)
{
    void *ptr = malloc(size);
    if (NULL == ptr) {
        abort();
    }

    return ptr;
}

LOCAL void checked_write(int fd, char *filename, void *buffer,
                         ssize_t count)
{
    ssize_t result = write(fd, buffer, count);
    if (-1 == result) {
        close(fd);
        croak("Could not write to the file %s: %s", filename,
              strerror(errno));
    }
    if (result != count) {
        close(fd);
        croak(
            "Write to %s did not write the expected amount of data (wrote %zd instead of %zu)",
            filename, result, count);
    }
}

LOCAL void checked_perlio_read(PerlIO * io, void *buffer,
                               SSize_t size)
{
    SSize_t read = PerlIO_read(io, buffer, size);
    check_perlio_result(read, size, "PerlIO_read");
}

LOCAL void check_perlio_result(SSize_t result, SSize_t expected,
                               char *op)
{
    if (result < 0) {
        croak("%s operation failed: %s\n", op, strerror(result));
    } else if (result != expected) {
        croak(
            "%s operation wrote %zd bytes when we expected to write %zd",
            op, result, expected);
    }
}
