#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <stdbool.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include "mmh3.h"
#include "bloom.h"

static const char *const bloom_errors[] = {
    "Success",
    "Out of memory",
    "Unable to open file",
    "Unable to read file",
    "Unable to write to file",
    "fstat() failure",
    "Invalid file format"
};

static const uint8_t bloom_magic[8] = { '!', 'b', 'l', 'o', 'o', 'm', 'f', '!' };

static bool
ideal_size(size_t expected, double accuracy, size_t *result)
{
    double bits = -((double)expected * log(accuracy)) / (log(2.0) * log(2.0));

    bits = ceil(bits);

    if (!isfinite(bits) || bits < 1.0 || bits > (double)SIZE_MAX) {
        return false;
    }

    *result = (size_t)bits;
    return true;
}


bloom_error_t bloom_init(bloomfilter *bf, const size_t expected, const double accuracy) {
    if (!bf || expected == 0 ||  !isfinite(accuracy) || accuracy <= 0.0 || accuracy >= 1.0) {
        return BF_INVALIDFILE;
    }

    float stored_accuracy = (float)accuracy;

    if (!isfinite(stored_accuracy) || stored_accuracy <= 0.0f || stored_accuracy >= 1.0f) {
        return BF_INVALIDFILE;
    }

    if (!ideal_size(expected, accuracy, &bf->size)) {
        return BF_INVALIDFILE;
    }

    bf->bitmap_size = bf->size / 8 + ((bf->size % 8) != 0);

    double hashes = ((double)bf->size / (double)expected) * log(2.0);

    if (!isfinite(hashes) || hashes < 1.0 || hashes > (double)BLOOM_MAX_HASHCOUNT) {
        return BF_INVALIDFILE;
    }

    bf->hashcount = (size_t)llround(hashes);

    bf->expected = expected;
    bf->accuracy = stored_accuracy;

    snprintf(bf->name, sizeof(bf->name), "DEFAULT");

    bf->bitmap = calloc(bf->bitmap_size, 1);

    if (!bf->bitmap) {
        return BF_OUTOFMEMORY;
    }

    return BF_SUCCESS;
}

void
bloom_destroy(bloomfilter *bf) {
    if (bf->bitmap) {
		free(bf->bitmap);
		bf->bitmap = NULL;
	}
}


// lookup table for bloom_saturation() and bloom_saturation_count()
static const uint8_t
bit_count_table[256] = {
    0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
    4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8
};


size_t
bloom_saturation_count(const bloomfilter *bf) {
	size_t count = 0;

	for (size_t i = 0; i < bf->bitmap_size; i++) {
		count += bit_count_table[bf->bitmap[i]];
	}

	return count;
}

float
bloom_saturation(const bloomfilter *bf) {
    size_t set_bits = bloom_saturation_count(bf);

    return (float)set_bits / (float)bf->size * 100.0f;
}


void
bloom_add(bloomfilter *bf, const void *element, const size_t len)
{
    uint64_t hash[2];

    mmh3_128(element, len, 0, hash);

    for (size_t i = 0; i < bf->hashcount; i++) {
        uint64_t result = (hash[0] + i * hash[1]) % bf->size;

        uint64_t byte_position = result / 8;
        uint8_t bit_position = result % 8;

        bf->bitmap[byte_position] |= (0x01 << bit_position);
    }
}

bool
bloom_set_name(bloomfilter *bf, const char *name)
{
    size_t len = strlen(name);

    if (!bloom_valid_filter_name((const uint8_t *)name, len + 1))
        return false;

    memcpy(bf->name, name, len + 1);
    return true;
}

static bloom_error_t
bloom_write_file(const bloomfilter *bf, FILE *fp)
{
    bloomfilter_file bff = {0};

    memcpy(bff.magic, bloom_magic, sizeof(bff.magic));

    bff.version     = BLOOM_FILE_VERSION;
    bff.size        = bf->size;
    bff.hashcount   = bf->hashcount;
    bff.bitmap_size = bf->bitmap_size;
    bff.expected    = bf->expected;
    bff.accuracy    = bf->accuracy;

    memcpy(bff.name, bf->name, sizeof(bff.name));

    if (fwrite(&bff, sizeof(bff), 1, fp) != 1) {
        return BF_FWRITE;
    }

    if (fwrite(bf->bitmap, 1, bf->bitmap_size, fp) != bf->bitmap_size) {
        return BF_FWRITE;
    }

    return BF_SUCCESS;
}

bloom_error_t
bloom_save(const bloomfilter *bf, const char *path)
{
    bloom_error_t result;
    FILE *fp = NULL;
    int fd = -1;

    size_t path_len = strlen(path);

    if (path_len > SIZE_MAX - sizeof(".tmp.XXXXXX")) {
        return BF_FOPEN;
    }

    char *tmp_path = malloc(path_len + sizeof(".tmp.XXXXXX"));
    if (tmp_path == NULL) {
        return BF_OUTOFMEMORY;
    }

    snprintf(tmp_path, path_len + sizeof(".tmp.XXXXXX"), "%s.tmp.XXXXXX", path);

    fd = mkstemp(tmp_path);
    if (fd < 0) {
        free(tmp_path);
        return BF_FOPEN;
    }

    fp = fdopen(fd, "wb");
    if (fp == NULL) {
        close(fd);
        unlink(tmp_path);
        free(tmp_path);
        return BF_FOPEN;
    }

    result = bloom_write_file(bf, fp);

    if (result != BF_SUCCESS)
        goto fail;

    if (fflush(fp) != 0) {
        result = BF_FWRITE;
        goto fail;
    }

    if (fsync(fd) != 0) {
        result = BF_FWRITE;
        goto fail;
    }

    if (fclose(fp) != 0) {
        fp = NULL;
        result = BF_FWRITE;
        goto fail_closed;
    }

    fp = NULL;

    if (rename(tmp_path, path) != 0) {
        result = BF_FWRITE;
        goto fail_closed;
    }

    free(tmp_path);
    return BF_SUCCESS;

fail:
    fclose(fp);
    fp = NULL;

fail_closed:
    unlink(tmp_path);
    free(tmp_path);

    return result;
}

const char *
bloom_strerror(bloom_error_t error) {
	if (error < 0 || error >= BF_ERRORCOUNT) {
		return "Unknown error";
	}

	return bloom_errors[error];
}

bool
bloom_valid_filter_name(const uint8_t *name, size_t size)
{
    size_t len = 0;

    while (len < size && name[len] != '\0') {
        len++;
    }

    if (len == 0 || len == size) {
        return false;
    }

    for (size_t i = 0; i < len; i++) {
        uint8_t c = name[i];

        if (!((c >= 'a' && c <= 'z') ||
              (c >= 'A' && c <= 'Z') ||
              (c >= '0' && c <= '9') ||
              c == ' ' || c == '-' || c == '_' || c == '.')) {
            return false;
        }
    }

    return true;
}

static bloom_error_t
validate_mapped_header(const bloomfilter_file *header, size_t file_size)
{
    if (memcmp(header->magic, bloom_magic, sizeof(header->magic)) != 0) {
        return BF_INVALIDFILE;
    }

    if (header->version != BLOOM_FILE_VERSION) return BF_INVALIDFILE;
    if (header->size == 0)                     return BF_INVALIDFILE;
    if (header->hashcount == 0)                return BF_INVALIDFILE;

    /*
     * Bound the amount of work for each lookup. Legitimate Bloom
     * filters should require nowhere near this many hashes.
     */
    if (header->hashcount > BLOOM_MAX_HASHCOUNT) {
        return BF_INVALIDFILE;
    }

    /*
     * Compute ceil(size / 8) without doing size + 7,
     * which could overflow.
     */
    uint64_t expected_bitmap_size = header->size / 8 + ((header->size % 8) != 0);
    if (header->bitmap_size != expected_bitmap_size) {
        return BF_INVALIDFILE;
    }

    /*
     * Make sure converting bitmap_size to size_t is safe.
     */
    if (header->bitmap_size > SIZE_MAX) {
        return BF_INVALIDFILE;
    }

    size_t bitmap_size = (size_t)header->bitmap_size;

    /*
     * Avoid overflow in sizeof(header) + bitmap_size
     */
    if (bitmap_size > SIZE_MAX - sizeof(bloomfilter_file)) {
        return BF_INVALIDFILE;
    }

    if (sizeof(bloomfilter_file) + bitmap_size != file_size) {
        return BF_INVALIDFILE;
    }

    /*
     * Validate filter name. Shouldn't contain many special
     * characters.
     */
    if (!bloom_valid_filter_name(header->name, sizeof(header->name))) {
      return BF_INVALIDFILE;
    }

    if (!isfinite(header->accuracy)) {
        return BF_INVALIDFILE;
    }

    if (header->accuracy <= 0.0f || header->accuracy >= 1.0f) {
        return BF_INVALIDFILE;
    }

    if (header->expected == 0) {
        return BF_INVALIDFILE;
    }

    return BF_SUCCESS;
}

bloom_error_t
bloom_map_fd(bloomfilter_mapped *bf, int fd)
{
    struct stat st;

    if (bf == NULL)                       return BF_INVALIDFILE;
    if (fstat(fd, &st) < 0)               return BF_FSTAT;
    if (st.st_size < 0)                   return BF_INVALIDFILE;
    if ((uintmax_t)st.st_size > SIZE_MAX) return BF_INVALIDFILE;

    size_t file_size = (size_t)st.st_size;
    if (file_size < sizeof(bloomfilter_file)) return BF_INVALIDFILE;

    void *mapping = mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
    if (mapping == MAP_FAILED) {
        return BF_FREAD;
    }

    const bloomfilter_file *header = (const bloomfilter_file *)mapping;

    bloom_error_t err = validate_mapped_header(header, file_size);
    if (err != BF_SUCCESS) {
        munmap(mapping, file_size);
        return err;
    }

    const uint8_t *bitmap = (const uint8_t *)mapping + sizeof(bloomfilter_file);

    bf->header = header;
    bf->bitmap = bitmap;
    bf->mapping = mapping;
    bf->mapping_size = file_size;

    return BF_SUCCESS;
}

void
bloom_unmap(bloomfilter_mapped *bf)
{
    if (bf->mapping) {
        munmap(bf->mapping, bf->mapping_size);
    }

    bf->header = NULL;
    bf->bitmap = NULL;
    bf->mapping = NULL;
    bf->mapping_size = 0;
}

bool
bloom_mapped_lookup(const bloomfilter_mapped *bf, const void *element, const size_t len)
{
    uint64_t hash[2];

    mmh3_128(element, len, 0, hash);

    for (size_t i = 0; i < bf->header->hashcount; i++) {
        uint64_t result = (hash[0] + i * hash[1]) % bf->header->size;

        uint64_t byte_position = result / 8;
        uint8_t bit_position = result % 8;

        if ((bf->bitmap[byte_position] & (0x01 << bit_position)) == 0) {
            return false;
        }
    }

    return true;
}