#include "stream.h"

#include <ropkg/compress.h>
#include <ropkg/stream.h>
#include <stdlib.h>
#include <string.h>

enum ropkg_status ropkg_stream_open(
	FILE *fp,
	const struct ropkg_compression_function *func,
	enum ropkg_stream_mode mode,
	struct ropkg_stream **out)
{
	switch (mode) {
	case ROPKG_STREAM_READ:
	case ROPKG_STREAM_WRITE:
		break;
	default:
		return ROPKG_ERR_INVALID_ARGUMENT;
	}

	struct ropkg_stream *stream = malloc(sizeof *stream);
	if (!stream) {
		return ROPKG_ERR_NO_MEMORY;
	}

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

	stream->s_fp = fp;
	stream->s_mode = mode;

	if (!func) {
		*out = stream;
		return ROPKG_SUCCESS;
	}

	struct ropkg_stream_compressor *c = &stream->s_compressor;
	enum ropkg_compression_stream_mode c_mode;
	switch (mode) {
	case ROPKG_STREAM_READ:
		c_mode = ROPKG_COMPRESSION_MODE_DECOMPRESS;
		break;
	case ROPKG_STREAM_WRITE:
		c_mode = ROPKG_COMPRESSION_MODE_COMPRESS;
		break;
	default:
		/* this won't happen */
		return ROPKG_ERR_INVALID_ARGUMENT;
	}

	enum ropkg_status status = ropkg_compression_function_get_buffer_size(
		func,
		c_mode,
		&c->c_in_max,
		&c->c_out_max);
	if (status != ROPKG_SUCCESS) {
		free(stream);
		return status;
	}

	c->c_in_buf = malloc(c->c_in_max);
	if (!c->c_in_buf) {
		free(stream);
		return ROPKG_ERR_NO_MEMORY;
	}

	c->c_out_buf = malloc(c->c_out_max);
	if (!c->c_out_buf) {
		free(c->c_in_buf);
		free(stream);
		return ROPKG_ERR_NO_MEMORY;
	}

	status = ropkg_compression_stream_open(
		func,
		c->c_in_buf,
		c->c_in_max,
		c->c_out_buf,
		c->c_out_max,
		c_mode,
		&c->c_stream);

	if (status != ROPKG_SUCCESS) {
		free(c->c_in_buf);
		free(c->c_out_buf);
		free(stream);
		return status;
	}

	c->c_func = func;

	*out = stream;
	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_close(struct ropkg_stream *stream)
{
	while (stream->s_compression_depth > 0) {
		ropkg_stream_end_compressed_section(stream, NULL, NULL);
	}

	if (stream->s_compressor.c_in_buf) {
		free(stream->s_compressor.c_in_buf);
	}

	if (stream->s_compressor.c_out_buf) {
		free(stream->s_compressor.c_out_buf);
	}

	if (stream->s_compressor.c_stream) {
		ropkg_compression_stream_close(stream->s_compressor.c_stream);
	}

	free(stream);

	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_begin_compressed_section(
	struct ropkg_stream *stream)
{
	if (stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED) {
		return ROPKG_ERR_BAD_STATE;
	}

	if (stream->s_compression_depth > 0) {
		stream->s_compression_depth++;
		return ROPKG_SUCCESS;
	}

	stream->s_compression_depth = 1;
	stream->s_compressor.c_in_count = 0;
	stream->s_compressor.c_nr_written_compressed = 0;
	stream->s_compressor.c_nr_written_uncompressed = 0;

	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_end_compressed_section(
	struct ropkg_stream *stream,
	size_t *out_nr_written_compressed,
	size_t *out_nr_written_uncompressed)
{
	if (stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED) {
		return ROPKG_ERR_BAD_STATE;
	}

	if (stream->s_compression_depth == 0) {
		return ROPKG_ERR_BAD_STATE;
	}

	if (stream->s_compression_depth > 1) {
		stream->s_compression_depth--;
		return ROPKG_SUCCESS;
	}

	stream->s_compression_depth = 0;
	struct ropkg_stream_compressor *c = &stream->s_compressor;
	if (out_nr_written_compressed) {
		*out_nr_written_compressed = c->c_nr_written_compressed;
	}

	if (out_nr_written_uncompressed) {
		*out_nr_written_uncompressed = c->c_nr_written_uncompressed;
	}

	if (c->c_nr_written_uncompressed == 0) {
		return ROPKG_SUCCESS;
	}

	size_t out_size = 0;

	ropkg_compression_stream_process(
		c->c_stream,
		c->c_in_count,
		ROPKG_COMPRESSION_OP_END,
		&out_size);

	size_t nr_written = fwrite(c->c_out_buf, 1, out_size, stream->s_fp);
	if (nr_written != out_size) {
		return ROPKG_ERR_IO_FAILURE;
	}

	c->c_nr_written_compressed += nr_written;
	stream->s_nr_written += out_size;
	if (out_nr_written_compressed) {
		*out_nr_written_compressed = c->c_nr_written_compressed;
	}

	stream->s_compressor.c_in_count = 0;
	stream->s_compressor.c_nr_written_compressed = 0;
	stream->s_compressor.c_nr_written_uncompressed = 0;

	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_read(
	struct ropkg_stream *stream,
	void *p,
	size_t len,
	size_t *nr_read)
{
	if (stream->s_mode != ROPKG_STREAM_READ) {
		return ROPKG_ERR_BAD_STATE;
	}

	return ROPKG_SUCCESS;
}

static enum ropkg_status flush_compressed_buffer(struct ropkg_stream *stream)
{
	struct ropkg_stream_compressor *c = &stream->s_compressor;

	size_t out_size;
	enum ropkg_status status = ropkg_compression_stream_process(
		c->c_stream,
		c->c_in_count,
		ROPKG_COMPRESSION_OP_CONTINUE,
		&out_size);

	if (status != ROPKG_SUCCESS) {
		return status;
	}

	size_t written = fwrite(c->c_out_buf, 1, out_size, stream->s_fp);

	c->c_nr_written_compressed += c->c_in_count;
	stream->s_nr_written_compressed = c->c_in_count;
	stream->s_nr_written += out_size;

	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_write(
	struct ropkg_stream *stream,
	const void *p,
	size_t len,
	size_t *nr_written)
{
	if (stream->s_mode != ROPKG_STREAM_WRITE) {
		return ROPKG_ERR_BAD_STATE;
	}

	size_t w = 0;
	if (stream->s_compression_depth == 0) {
		w = fwrite(p, 1, len, stream->s_fp);

		*nr_written = w;
		if (w != len) {
			return ROPKG_ERR_IO_FAILURE;
		}

		if (!(stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED)) {
			stream->s_nr_written_uncompressed += w;
			stream->s_nr_written += w;
		}

		return ROPKG_SUCCESS;
	}

	if (stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED) {
		return ROPKG_ERR_BAD_STATE;
	}

	struct ropkg_stream_compressor *c = &stream->s_compressor;
	const unsigned char *src = p;
	enum ropkg_status status = ROPKG_SUCCESS;

	size_t remaining = len;
	while (remaining > 0 && status == ROPKG_SUCCESS) {
		unsigned char *dest = c->c_in_buf + c->c_in_count;
		size_t space_remaining = c->c_in_max - c->c_in_count;

		size_t to_copy = remaining;
		if (to_copy > space_remaining) {
			to_copy = space_remaining;
		}

		memcpy(dest, src, to_copy);

		c->c_in_count += to_copy;
		c->c_nr_written_uncompressed += to_copy;
		stream->s_nr_written_uncompressed += to_copy;
		space_remaining -= to_copy;
		remaining -= to_copy;
		w += to_copy;

		if (space_remaining == 0) {
			status = flush_compressed_buffer(stream);
		}
	}

	*nr_written = w;
	return status;
}

size_t ropkg_stream_get_cursor_position(struct ropkg_stream *stream)
{
	return ftell(stream->s_fp);
}

enum ropkg_status ropkg_stream_set_cursor_position(
	struct ropkg_stream *stream,
	size_t pos)
{
	if (stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED) {
		return ROPKG_ERR_BAD_STATE;
	}

	if (stream->s_compression_depth > 0) {
		return ROPKG_ERR_BAD_STATE;
	}

	stream->s_saved_cursor = ftell(stream->s_fp);
	fseek(stream->s_fp, pos, SEEK_SET);

	stream->s_flags |= ROPKG_STREAM_F_CURSOR_MOVED;

	return ROPKG_SUCCESS;
}

enum ropkg_status ropkg_stream_restore_cursor_position(
	struct ropkg_stream *stream)
{
	if (!(stream->s_flags & ROPKG_STREAM_F_CURSOR_MOVED)) {
		return ROPKG_ERR_BAD_STATE;
	}

	fseek(stream->s_fp, stream->s_saved_cursor, SEEK_SET);
	stream->s_saved_cursor = 0;

	stream->s_flags &= ~ROPKG_STREAM_F_CURSOR_MOVED;

	return ROPKG_SUCCESS;
}

bool ropkg_stream_is_in_compressed_section(struct ropkg_stream *stream)
{
	return stream->s_compression_depth > 0;
}

size_t ropkg_stream_get_nr_written(struct ropkg_stream *stream)
{
	return stream->s_nr_written;
}

size_t ropkg_stream_get_nr_written_compressed(struct ropkg_stream *stream)
{
	return stream->s_nr_written_compressed;
}

size_t ropkg_stream_get_nr_written_uncompressed(struct ropkg_stream *stream)
{
	return stream->s_nr_written_uncompressed;
}