#include <blue/compress/zstd.h>
#include <blue/core/ringbuffer.h>
#include <zstd.h>

/*** PRIVATE DATA *************************************************************/

struct b_zstd_compressor_p {
	union {
		ZSTD_CCtx *zstd_c;
		ZSTD_DCtx *zstd_d;
	};
};

/*** PUBLIC FUNCTIONS *********************************************************/

b_status b_zstd_compressor_get_buffer_size(
	b_compressor_mode mode, size_t *inbuf_size, size_t *outbuf_size)
{
	switch (mode) {
	case B_COMPRESSOR_MODE_COMPRESS:
		*inbuf_size = ZSTD_CStreamInSize();
		*outbuf_size = ZSTD_CStreamOutSize();
		break;
	case B_COMPRESSOR_MODE_DECOMPRESS:
		*inbuf_size = ZSTD_DStreamInSize();
		*outbuf_size = ZSTD_DStreamOutSize();
		break;
	default:
		return B_ERR_INVALID_ARGUMENT;
	}

	return B_SUCCESS;
}

/*** VIRTUAL FUNCTIONS ********************************************************/

static void zstd_compressor_init(b_object *obj, void *priv)
{
}

static void zstd_compressor_fini(b_object *obj, void *priv)
{
	b_compressor_data *c = b_object_get_protected(obj, B_TYPE_COMPRESSOR);
	struct b_zstd_compressor_p *ctx = priv;
	switch (c->c_mode) {
	case B_COMPRESSOR_MODE_COMPRESS:
		ZSTD_freeCCtx(ctx->zstd_c);
		break;
	case B_COMPRESSOR_MODE_DECOMPRESS:
		ZSTD_freeDCtx(ctx->zstd_d);
		break;
	default:
		break;
	}
}

static enum b_status reset(b_compressor *compressor)
{
	struct b_zstd_compressor_p *ctx
		= b_object_get_private(compressor, B_TYPE_ZSTD_COMPRESSOR);
	b_compressor_data *data
		= b_object_get_protected(compressor, B_TYPE_COMPRESSOR);

	if (!ctx || !data) {
		return B_ERR_INVALID_ARGUMENT;
	}

	switch (data->c_mode) {
	case B_COMPRESSOR_MODE_COMPRESS:
		ZSTD_CCtx_reset(ctx->zstd_c, ZSTD_reset_session_only);
		break;
	case B_COMPRESSOR_MODE_DECOMPRESS:
		ZSTD_DCtx_reset(ctx->zstd_d, ZSTD_reset_session_only);
		break;
	default:
		return B_ERR_BAD_STATE;
	}

	return B_SUCCESS;
}

static enum b_status compress(b_compressor *compressor)
{
	enum b_status status = B_SUCCESS;
	struct b_zstd_compressor_p *ctx
		= b_object_get_private(compressor, B_TYPE_ZSTD_COMPRESSOR);
	b_compressor_data *data
		= b_object_get_protected(compressor, B_TYPE_COMPRESSOR);

	if (!ctx || !data) {
		return B_ERR_INVALID_ARGUMENT;
	}

	b_ringbuffer *in = data->c_in;
	b_ringbuffer *out = data->c_out;

	if (b_ringbuffer_available_data_remaining(in) == 0) {
		return B_ERR_NO_DATA;
	}

	if (b_ringbuffer_write_capacity_remaining(out) == 0) {
		return B_ERR_NO_SPACE;
	}

	size_t nr_consumed = 0;

	while (1) {
		size_t in_available = 0, out_capacity = 0;
		const void *in_buf = NULL;
		void *out_buf = NULL;

		status = b_ringbuffer_open_read_buffer(in, &in_buf, &in_available);
		if (!B_OK(status)) {
			break;
		}

		status = b_ringbuffer_open_write_buffer(
			out, &out_buf, &out_capacity);
		if (!B_OK(status)) {
			b_ringbuffer_close_read_buffer(in, &in_buf, 0);
			break;
		}

		ZSTD_inBuffer z_in = {
			.src = in_buf,
			.pos = 0,
			.size = in_available,
		};

		ZSTD_outBuffer z_out = {
			.dst = out_buf,
			.pos = 0,
			.size = out_capacity,
		};

		do {
			size_t ret = ZSTD_compressStream2(
				ctx->zstd_c, &z_out, &z_in, ZSTD_e_continue);
			if (ZSTD_isError(ret)) {
				status = B_ERR_COMPRESSION_FAILURE;
				break;
			}
		} while (z_in.pos < z_in.size && z_out.pos < z_out.size);

		nr_consumed += z_in.pos;

		b_ringbuffer_close_read_buffer(in, &in_buf, z_in.pos);
		b_ringbuffer_close_write_buffer(out, &out_buf, z_out.pos);
	}

	if ((status == B_ERR_NO_SPACE || status == B_ERR_NO_DATA)
	    && nr_consumed > 0) {
		status = B_SUCCESS;
	}

	return status;
}

static enum b_status compress_end(b_compressor *compressor)
{
	enum b_status status = B_SUCCESS;
	struct b_zstd_compressor_p *ctx
		= b_object_get_private(compressor, B_TYPE_ZSTD_COMPRESSOR);
	b_compressor_data *data
		= b_object_get_protected(compressor, B_TYPE_COMPRESSOR);

	if (!ctx || !data) {
		return B_ERR_INVALID_ARGUMENT;
	}

	b_ringbuffer *out = data->c_out;
	if (b_ringbuffer_write_capacity_remaining(out) == 0) {
		return B_ERR_NO_SPACE;
	}

	bool finished = false;
	do {
		void *out_buf = NULL;
		size_t out_capacity = 0;
		status = b_ringbuffer_open_write_buffer(
			out, &out_buf, &out_capacity);
		if (!B_OK(status)) {
			break;
		}

		ZSTD_inBuffer z_in = {0};
		ZSTD_outBuffer z_out = {
			.dst = out_buf,
			.pos = 0,
			.size = out_capacity,
		};

		do {
			size_t ret = ZSTD_compressStream2(
				ctx->zstd_c, &z_out, &z_in, ZSTD_e_end);
			if (ZSTD_isError(ret)) {
				status = B_ERR_COMPRESSION_FAILURE;
				finished = true;
			}

			if (ret == 0) {
				data->c_flags |= B_COMPRESSOR_EOF;
				finished = true;
			}
		} while (!finished && z_out.pos < z_out.size);

		b_ringbuffer_close_write_buffer(out, &out_buf, z_out.pos);
	} while (!finished);

	return status;
}

static enum b_status decompress(b_compressor *compressor)
{
	enum b_status status = B_SUCCESS;
	struct b_zstd_compressor_p *ctx
		= b_object_get_private(compressor, B_TYPE_ZSTD_COMPRESSOR);
	b_compressor_data *data
		= b_object_get_protected(compressor, B_TYPE_COMPRESSOR);

	if (!ctx || !data) {
		return B_ERR_INVALID_ARGUMENT;
	}

	b_ringbuffer *in = data->c_in;
	b_ringbuffer *out = data->c_out;

	if (b_ringbuffer_available_data_remaining(in) == 0) {
		return B_ERR_NO_DATA;
	}

	if (b_ringbuffer_write_capacity_remaining(out) == 0) {
		return B_ERR_NO_SPACE;
	}

	size_t nr_consumed = 0;

	while (!(data->c_flags & B_COMPRESSOR_EOF)) {
		size_t in_available = 0, out_capacity = 0;
		const void *in_buf = NULL;
		void *out_buf = NULL;

		status = b_ringbuffer_open_read_buffer(in, &in_buf, &in_available);
		if (!B_OK(status)) {
			break;
		}

		status = b_ringbuffer_open_write_buffer(
			out, &out_buf, &out_capacity);
		if (!B_OK(status)) {
			b_ringbuffer_close_read_buffer(in, &in_buf, 0);
			break;
		}

		ZSTD_inBuffer z_in = {
			.src = in_buf,
			.pos = 0,
			.size = in_available,
		};

		ZSTD_outBuffer z_out = {
			.dst = out_buf,
			.pos = 0,
			.size = out_capacity,
		};

		do {
			size_t ret = ZSTD_decompressStream(
				ctx->zstd_d, &z_out, &z_in);
			if (ZSTD_isError(ret)) {
				status = B_ERR_COMPRESSION_FAILURE;
				break;
			}

			if (ret == 0) {
				data->c_flags |= B_COMPRESSOR_EOF;
				break;
			}
		} while (z_in.pos < z_in.size && z_out.pos < z_out.size);

		nr_consumed += z_in.pos;

		b_ringbuffer_close_read_buffer(in, &in_buf, z_in.pos);
		b_ringbuffer_close_write_buffer(out, &out_buf, z_out.pos);
	}

	if ((status == B_ERR_NO_SPACE || status == B_ERR_NO_DATA)
	    && nr_consumed > 0) {
		status = B_SUCCESS;
	}

	return status;
}

static enum b_status set_mode(b_compressor *compressor, b_compressor_mode mode)
{
	struct b_zstd_compressor_p *ctx
		= b_object_get_private(compressor, B_TYPE_ZSTD_COMPRESSOR);
	b_compressor_data *data
		= b_object_get_protected(compressor, B_TYPE_COMPRESSOR);

	if (!ctx || !data) {
		return B_ERR_INVALID_ARGUMENT;
	}

	if (mode == data->c_mode) {
		return B_SUCCESS;
	}

	switch (data->c_mode) {
	case B_COMPRESSOR_MODE_COMPRESS:
		ZSTD_freeCCtx(ctx->zstd_c);
		break;
	case B_COMPRESSOR_MODE_DECOMPRESS:
		ZSTD_freeDCtx(ctx->zstd_d);
		break;
	default:
		break;
	}

	data->c_mode = mode;

	switch (data->c_mode) {
	case B_COMPRESSOR_MODE_COMPRESS:
		ctx->zstd_c = ZSTD_createCCtx();
		break;
	case B_COMPRESSOR_MODE_DECOMPRESS:
		ctx->zstd_d = ZSTD_createDCtx();
		break;
	default:
		return B_ERR_INVALID_ARGUMENT;
	}

	return B_SUCCESS;
}

/*** CLASS DEFINITION *********************************************************/

B_TYPE_CLASS_DEFINITION_BEGIN(b_zstd_compressor)
	B_TYPE_CLASS_INTERFACE_BEGIN(b_object, B_TYPE_OBJECT)
		B_INTERFACE_ENTRY(to_string) = NULL;
	B_TYPE_CLASS_INTERFACE_END(b_object, B_TYPE_OBJECT)

	B_TYPE_CLASS_INTERFACE_BEGIN(b_compressor, B_TYPE_COMPRESSOR)
		B_INTERFACE_ENTRY(c_buffer_size)
			= b_zstd_compressor_get_buffer_size;
		B_INTERFACE_ENTRY(c_compress) = compress;
		B_INTERFACE_ENTRY(c_compress_end) = compress_end;
		B_INTERFACE_ENTRY(c_decompress) = decompress;
		B_INTERFACE_ENTRY(c_reset) = reset;
		B_INTERFACE_ENTRY(c_set_mode) = set_mode;
	B_TYPE_CLASS_INTERFACE_END(b_compressor, B_TYPE_COMPRESSOR)
B_TYPE_CLASS_DEFINITION_END(b_zstd_compressor)

B_TYPE_DEFINITION_BEGIN(b_zstd_compressor)
	B_TYPE_ID(0x51d437fc, 0xe789, 0x4105, 0xbac7, 0xe6b3f45df198);
	B_TYPE_EXTENDS(B_TYPE_COMPRESSOR);
	B_TYPE_CLASS(b_zstd_compressor_class);
	B_TYPE_INSTANCE_PRIVATE(struct b_zstd_compressor_p);
	B_TYPE_INSTANCE_INIT(zstd_compressor_init);
	B_TYPE_INSTANCE_FINI(zstd_compressor_fini);
B_TYPE_DEFINITION_END(b_zstd_compressor)