#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)) {
        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);

    SV *merged_ref = newRV_inc((SV *)hash_new);
    /* We aren't keeping the original HV * around so we decrement its ref
     * count. */
    SvREFCNT_dec(hash_new);

    return merged_ref;
}

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;
            }

            if (next_is_right) {
                record->type = MMDBW_RECORD_TYPE_NODE;
                record->value.node = next_node;
            } else {
                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!");
    }

    SvREFCNT_inc_simple_void_NN(thawed);

    PUTBACK;
    FREETMPS;
    LEAVE;

    return (HV *)SvRV(thawed);
}

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);
    }
}
