#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include <bzlib.h>
#include <erl_nif.h>

// TODO: Parameterize.  This has been tuned so that the decompression NIF calls
// return in less than 1ms on a developer's machine, but mileage may vary, and
// compression hasn't been tuned at all.
static size_t default_output_buffer_size = 10 * 1024;

// Uncomment for verbose output.
// #define BZIP2_EX_DEBUG

#ifdef BZIP2_EX_DEBUG
static int verbosity = 4;

static void debug_print(const char* fmt, ...) {
	va_list args;
	va_start(args, fmt);
	vfprintf(stderr, fmt, args);
	va_end(args);
}
#else
static int verbosity = 0;
static inline void debug_print(const char* fmt, ...) {}
#endif

static void* bzalloc_erl(void* opaque, int num, int unit_size) {
	return enif_alloc(num * unit_size);
}

void bzfree_erl(void* opaque, void* ptr) {
	enif_free(ptr);
}

// An opaque structure passed back to elixir, to support continuation.
static ErlNifResourceType *erl_bz_stream_wrapper;

typedef struct {
	void* input_buffer;
	void* output_buffer;
	bz_stream* stream;
} bz_stream_wrapper;

static ERL_NIF_TERM atom_ok;
static ERL_NIF_TERM atom_cont;
static ERL_NIF_TERM atom_next_input;
static ERL_NIF_TERM atom_stream_end;
static ERL_NIF_TERM atom_bad_magic;
static ERL_NIF_TERM atom_sequence_error;
static ERL_NIF_TERM atom_param_error;
static ERL_NIF_TERM atom_mem_error;
static ERL_NIF_TERM atom_data_error;
static ERL_NIF_TERM atom_io_error;
static ERL_NIF_TERM atom_unexpected_eof;
static ERL_NIF_TERM atom_outbuff_full;
static ERL_NIF_TERM atom_config_error;
static ERL_NIF_TERM atom_unknown_error;

static ERL_NIF_TERM
convert_error_to_atom(int bz2_error) {
	switch (bz2_error) {
		case BZ_SEQUENCE_ERROR: return atom_sequence_error;
		case BZ_PARAM_ERROR: return atom_param_error;
		case BZ_MEM_ERROR: return atom_mem_error;
		case BZ_DATA_ERROR: return atom_data_error;
		case BZ_IO_ERROR: return atom_io_error;
		case BZ_UNEXPECTED_EOF: return atom_unexpected_eof;
		case BZ_OUTBUFF_FULL: return atom_outbuff_full;
		case BZ_CONFIG_ERROR: return atom_config_error;
		default: return atom_unknown_error;
	}
}

static void
destroy_stream_wrapper(ErlNifEnv *env, void *res) {
	bz_stream_wrapper** resource_ptr = res;
	bz_stream_wrapper* resource = *resource_ptr;
	if (resource != NULL) {
		if (resource->input_buffer != NULL) {
			enif_free(resource->input_buffer);
		}
		if (resource->output_buffer != NULL) {
			enif_free(resource->output_buffer);
		}
		if (resource->stream != NULL) {
			enif_free(resource->stream);
		}
	}
}

static bz_stream_wrapper*
allocate_stream_wrapper(ErlNifEnv *env) {
	bz_stream_wrapper* state = enif_alloc_resource(erl_bz_stream_wrapper, sizeof(bz_stream_wrapper));
	state->stream = enif_alloc(sizeof(bz_stream));
	if (state->stream == NULL) {
		destroy_stream_wrapper(env, &state);
		return NULL;
	}
	state->input_buffer = NULL;
	state->output_buffer = enif_alloc(default_output_buffer_size);
	if (state->output_buffer == NULL) {
		destroy_stream_wrapper(env, &state);
		return NULL;
	}
	state->stream->bzalloc = &bzalloc_erl;
	state->stream->bzfree = &bzfree_erl;
	state->stream->opaque = NULL;
	return state;
}

static ERL_NIF_TERM
compress_init(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state = allocate_stream_wrapper(env);
	if (state == NULL) {
		return enif_raise_exception(env, atom_mem_error);
	}

	// TODO: parameterize
	int blockSize100k = 9;
	int defaultWorkFactor = 0;
	int result = BZ2_bzCompressInit(state->stream, blockSize100k, verbosity, defaultWorkFactor);
	if (result != BZ_OK) {
		destroy_stream_wrapper(env, &state);
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	return enif_make_resource(env, state);
}

static ERL_NIF_TERM
compress_block(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state;
	if (!enif_get_resource(env, argv[0], erl_bz_stream_wrapper, (void**)&state)) {
		return enif_make_badarg(env);
	}
	bz_stream* stream = state->stream;

	if (state->input_buffer == NULL) {
		// Copy NIF binary input over to the bz_stream and continuation state.
		ErlNifBinary input_bin;
		if (!enif_inspect_iolist_as_binary(env, argv[1], &input_bin)) {
			return enif_make_badarg(env);
		}
		state->input_buffer = enif_alloc(input_bin.size);
		if (state->input_buffer == NULL) {
			return enif_raise_exception(env, atom_mem_error);
		}
		stream->avail_in = input_bin.size;
		stream->next_in = state->input_buffer;
		memcpy(stream->next_in, input_bin.data, input_bin.size);
	} else {
		// Continuing to read from the input (last run was limited by the output
		// buffer size)
	}

	// Rewind output cursor for every run, we've already copied the previous run's
	// data back to the beam.
	stream->avail_out = default_output_buffer_size;
	stream->next_out = state->output_buffer;

	// FIXME: input "" is magic to put into finishing mode.
	int action = stream->avail_in ? BZ_RUN : BZ_FINISH;
	debug_print("compressing: %u bytes with action %d\n", stream->avail_in, action);
	int result = BZ2_bzCompress(stream, action);
	debug_print("compressed: out_bytes=%u\n", default_output_buffer_size - stream->avail_out);

	ERL_NIF_TERM status;
	if (result == BZ_RUN_OK || result == BZ_FINISH_OK) {
		// Ran out of input or output, continue working on this cached input
		// buffer but switch to finishing mode.
		status = atom_cont;
	} else if (result == BZ_STREAM_END) {
		// Finished with finishing.
		status = atom_stream_end;
	} else {
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	size_t bytes_written = default_output_buffer_size - stream->avail_out;
	ERL_NIF_TERM output_term;
	unsigned char* output_bin = enif_make_new_binary(env, bytes_written, &output_term);
	memcpy(output_bin, state->output_buffer, bytes_written);

	return enif_make_tuple3(env, atom_ok, status, output_term);
}

static ERL_NIF_TERM
compress_end(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state;
	if (!enif_get_resource(env, argv[0], erl_bz_stream_wrapper, (void**)&state)) {
		return enif_make_badarg(env);
	}

	int result = BZ2_bzCompressEnd(state->stream);
	if (result != BZ_OK) {
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	return atom_ok;
}

static ERL_NIF_TERM
decompress_init(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state = allocate_stream_wrapper(env);
	if (state == NULL) {
		return enif_raise_exception(env, atom_mem_error);
	}

	int small_memory = 0;
	int result = BZ2_bzDecompressInit(state->stream, verbosity, small_memory);
	if (result != BZ_OK) {
		destroy_stream_wrapper(env, &state);
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	return enif_make_resource(env, state);
}

static ERL_NIF_TERM
decompress_block(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state;
	if (!enif_get_resource(env, argv[0], erl_bz_stream_wrapper, (void**)&state)) {
		return enif_make_badarg(env);
	}
	bz_stream* stream = state->stream;

	// Clean up from the previous block.
	if (stream->avail_in == 0 && state->input_buffer != NULL) {
		enif_free(state->input_buffer);
		state->input_buffer = NULL;
	}
	// Starting with new input.
	if (state->input_buffer == NULL) {
		// Copy NIF binary input over to the bz_stream and continuation state.
		ErlNifBinary input_bin;
		if (!enif_inspect_iolist_as_binary(env, argv[1], &input_bin)) {
			return enif_make_badarg(env);
		}
		state->input_buffer = enif_alloc(input_bin.size);
		if (state->input_buffer == NULL) {
			return enif_raise_exception(env, atom_mem_error);
		}
		stream->avail_in = input_bin.size;
		stream->next_in = state->input_buffer;
		memcpy(stream->next_in, input_bin.data, input_bin.size);
	} else {
		// Continuing to read from the input (last run was limited by the output
		// buffer size)
	}

	stream->avail_out = default_output_buffer_size;
	stream->next_out = state->output_buffer;

	int result = BZ2_bzDecompress(stream);

	ERL_NIF_TERM status;
	if (result == BZ_OK) {
		if (stream->avail_in == 0) {
			// Ran out of input, continue with the next block.
			status = atom_next_input;
		} else if (stream->avail_out == 0) {
			// Ran out of output buffer space, return and continue with the cached
			// input.
			status = atom_cont;
		} else {
			// Should be unreachable.
			return enif_raise_exception(env, atom_unknown_error);
		}
	} else if (result == BZ_STREAM_END) {
		status = enif_make_tuple2(env,
			atom_stream_end,
			enif_make_int(env, stream->avail_in));
	} else if (result == BZ_DATA_ERROR_MAGIC) {
		debug_print("Ignoring trailing garbage after EOF.\n");
		status = atom_bad_magic;
	} else {
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	size_t bytes_written = default_output_buffer_size - stream->avail_out;
	ERL_NIF_TERM output_term;
	unsigned char* output_bin = enif_make_new_binary(env, bytes_written, &output_term);
	memcpy(output_bin, state->output_buffer, bytes_written);

	return enif_make_tuple3(env, atom_ok, status, output_term);
}

static ERL_NIF_TERM
decompress_end(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
	bz_stream_wrapper* state;
	if (!enif_get_resource(env, argv[0], erl_bz_stream_wrapper, (void**)&state)) {
		return enif_make_badarg(env);
	}

	int result = BZ2_bzDecompressEnd(state->stream);
	if (result != BZ_OK) {
		return enif_raise_exception(env, convert_error_to_atom(result));
	}

	return atom_ok;
}

static int
load(ErlNifEnv *env, void **priv, ERL_NIF_TERM info) {
	erl_bz_stream_wrapper = enif_open_resource_type(env, NULL, "bzip2_stream_wrapper",
		&destroy_stream_wrapper, ERL_NIF_RT_CREATE, NULL);
	if (erl_bz_stream_wrapper == NULL) {
		return -1;
	}

	atom_ok = enif_make_atom(env, "ok");
	atom_cont = enif_make_atom(env, "cont");
	atom_next_input = enif_make_atom(env, "next_input");
	atom_stream_end = enif_make_atom(env, "stream_end");
	atom_bad_magic = enif_make_atom(env, "bad_magic");
	atom_sequence_error = enif_make_atom(env, "bz2_sequence_error");
	atom_param_error = enif_make_atom(env, "bz2_param_error");
	atom_mem_error = enif_make_atom(env, "bz2_mem_error");
	atom_data_error = enif_make_atom(env, "bz2_data_error");
	atom_io_error = enif_make_atom(env, "bz2_io_error");
	atom_unexpected_eof = enif_make_atom(env, "bz2_unexpected_eof");
	atom_outbuff_full = enif_make_atom(env, "bz2_outbuff_full");
	atom_config_error = enif_make_atom(env, "bz2_config_error");
	atom_unknown_error = enif_make_atom(env, "bz2_unknown_error");

	return 0;
}

static ErlNifFunc funcs[] = {
	{"compress_init!", 0, &compress_init, 0},
	{"compress_block!", 2, &compress_block, 0},
	{"compress_end!", 1, &compress_end, 0},
	{"decompress_init!", 0, &decompress_init, 0},
	{"decompress_block!", 2, &decompress_block, 0},
	{"decompress_end!", 1, &decompress_end, 0}
};

ERL_NIF_INIT(Elixir.Bzip2.Driver, funcs, &load, NULL, NULL, NULL)