#include "compress.h"

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

struct zstd_stream {
	struct ropkg_compression_stream s_base;
	union {
		ZSTD_CCtx *c;
		ZSTD_DCtx *d;
	} s_ctx;
};

static enum ropkg_status zstd_buffer_size(
	enum ropkg_compression_stream_mode mode,
	size_t *in_buffer_size,
	size_t *out_buffer_size)
{
	switch (mode) {
	case ROPKG_COMPRESSION_MODE_COMPRESS:
		if (in_buffer_size) {
			*in_buffer_size = ZSTD_CStreamInSize();
		}

		if (out_buffer_size) {
			*out_buffer_size = ZSTD_CStreamOutSize();
		}
		return ROPKG_SUCCESS;
	case ROPKG_COMPRESSION_MODE_DECOMPRESS:
		if (in_buffer_size) {
			*in_buffer_size = ZSTD_DStreamInSize();
		}

		if (out_buffer_size) {
			*out_buffer_size = ZSTD_DStreamOutSize();
		}
		return ROPKG_SUCCESS;
	default:
		return ROPKG_ERR_INVALID_ARGUMENT;
	}
}

static enum ropkg_status zstd_open(
	const struct ropkg_compression_function *func,
	enum ropkg_compression_stream_mode mode,
	void *in_buffer,
	size_t in_max,
	void *out_buffer,
	size_t out_max,
	struct ropkg_compression_stream **out)
{
	struct zstd_stream *stream = malloc(sizeof *stream);
	if (!stream) {
		return ROPKG_ERR_NO_MEMORY;
	}

	memset(stream, 0x0, sizeof *stream);

	stream->s_base.s_in = in_buffer;
	stream->s_base.s_in_max = in_max;
	stream->s_base.s_out = out_buffer;
	stream->s_base.s_out_max = out_max;
	stream->s_base.s_func = func;
	stream->s_base.s_mode = mode;

	switch (mode) {
	case ROPKG_COMPRESSION_MODE_COMPRESS:
		stream->s_ctx.c = ZSTD_createCCtx();

		if (!stream->s_ctx.c) {
			free(stream);
			return ROPKG_ERR_NO_MEMORY;
		}

		ZSTD_CCtx_setParameter(
			stream->s_ctx.c,
			ZSTD_c_compressionLevel,
			10);
		ZSTD_CCtx_setParameter(stream->s_ctx.c, ZSTD_c_checksumFlag, 1);
		break;
	case ROPKG_COMPRESSION_MODE_DECOMPRESS:
		stream->s_ctx.d = ZSTD_createDCtx();

		if (!stream->s_ctx.d) {
			free(stream);
			return ROPKG_ERR_NO_MEMORY;
		}

		break;
	default:
		free(stream);
		return ROPKG_ERR_INVALID_ARGUMENT;
	}

	*out = &stream->s_base;
	return ROPKG_SUCCESS;
}

static enum ropkg_status zstd_close(struct ropkg_compression_stream *stream)
{
	struct zstd_stream *zstream = (struct zstd_stream *)stream;
	switch (stream->s_mode) {
	case ROPKG_COMPRESSION_MODE_COMPRESS:
		ZSTD_freeCCtx(zstream->s_ctx.c);
		break;
	case ROPKG_COMPRESSION_MODE_DECOMPRESS:
		ZSTD_freeDCtx(zstream->s_ctx.d);
		break;
	default:
		return ROPKG_ERR_BAD_STATE;
	}

	free(stream);
	return ROPKG_SUCCESS;
}

static enum ropkg_status zstd_compress(
	struct zstd_stream *zstream,
	enum ropkg_compression_stream_op op,
	ZSTD_inBuffer *in,
	ZSTD_outBuffer *out)
{
	ZSTD_EndDirective zop;

	switch (op) {
	case ROPKG_COMPRESSION_OP_CONTINUE:
		zop = ZSTD_e_continue;
		break;
	case ROPKG_COMPRESSION_OP_END:
		zop = ZSTD_e_end;
		break;
	default:
		return ROPKG_ERR_INVALID_ARGUMENT;
	}

	do {
		size_t ret
			= ZSTD_compressStream2(zstream->s_ctx.c, out, in, zop);

		if (ret == 0) {
			break;
		}
	} while (in->pos < in->size);

	return ROPKG_SUCCESS;
}

static enum ropkg_status zstd_decompress(
	struct zstd_stream *zstream,
	ZSTD_inBuffer *in,
	ZSTD_outBuffer *out)
{
	size_t ret = 0;

	while (in->pos < in->size) {
		ret = ZSTD_decompressStream(zstream->s_ctx.d, out, in);
	}

	if (ret != 0) {
		return ROPKG_ERR_INTERNAL_FAILURE;
	}

	return ROPKG_SUCCESS;
}

static enum ropkg_status zstd_process(
	struct ropkg_compression_stream *stream,
	size_t in_size,
	enum ropkg_compression_stream_op op,
	size_t *out_size)
{
	struct zstd_stream *zstream = (struct zstd_stream *)stream;
	ZSTD_inBuffer in = {.src = stream->s_in, .size = in_size};
	ZSTD_outBuffer out = {.dst = stream->s_out, .size = stream->s_out_max};

	enum ropkg_status status;
	switch (stream->s_mode) {
	case ROPKG_COMPRESSION_MODE_COMPRESS:
		status = zstd_compress(zstream, op, &in, &out);
		break;
	case ROPKG_COMPRESSION_MODE_DECOMPRESS:
		status = zstd_decompress(zstream, &in, &out);
		break;
	default:
		return ROPKG_ERR_BAD_STATE;
	}

	if (out_size) {
		*out_size = out.pos;
	}

	return ROPKG_SUCCESS;
}

const struct ropkg_compression_function ropkg_zstd = {
	.f_name = "zstd",
	.f_extension = ".zst",
	.f_buffer_size = zstd_buffer_size,
	.f_open = zstd_open,
	.f_close = zstd_close,
	.f_process = zstd_process,
};