#include <zlib.h>
#include <lauxlib.h>
#include <lua.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>

#include "lib.h"

#define DECOMPRESS_META "zlib.decompress.meta"

typedef struct {
	z_stream stream;
	size_t processed;
	uint8_t buff[4096];
	bool finalized;
} lua_zlib_stream_t;

static int lib_decompress_iter(lua_State *ctx) {
	size_t size;
	uint8_t *data = (uint8_t*)lua_tolstring(ctx, lua_upvalueindex(1), &size);

	lua_zlib_stream_t *stream = lua_touserdata(ctx, lua_upvalueindex(2));

	if (stream->finalized) {
		lua_pushnil(ctx);
		return 1;
	}

	stream->stream.next_in = data + stream->processed;
	stream->stream.avail_in = size - stream->processed;

	stream->stream.next_out = stream->buff;
	stream->stream.avail_out = sizeof stream->buff;

	int res = inflate(&stream->stream, Z_NO_FLUSH);

	switch (res) {
		case Z_STREAM_END:
			inflateEnd(&stream->stream);
			stream->finalized = true;
		case Z_OK: {
			lua_pushlstring(ctx, (const char*)stream->buff, sizeof stream->buff - stream->stream.avail_out);
			stream->processed = size - stream->stream.avail_in;
			return 1;
		}
		case Z_BUF_ERROR:
			inflateEnd(&stream->stream);
			stream->finalized = true;
			return luaL_error(ctx, "zlib error: incorrect C parameters passed");
		case Z_MEM_ERROR:
			inflateEnd(&stream->stream);
			stream->finalized = true;
			return luaL_error(ctx, "zlib error: memory ran out");
		case Z_STREAM_ERROR:
			inflateEnd(&stream->stream);
			stream->finalized = true;
			return luaL_error(ctx, "zlib error: invalid state");
		case Z_DATA_ERROR:
			inflateEnd(&stream->stream);
			stream->finalized = true;
			return luaL_error(ctx, "zlib error: %s", stream->stream.msg);
		default:
			return luaL_error(ctx, "wtf");
	}
}
static int lib_decompress_free(lua_State *ctx) {
	lua_zlib_stream_t *stream = lua_touserdata(ctx, 1);
	if (stream->finalized) return 0;

	inflateEnd(&stream->stream);
	stream->finalized = true;

	return 0;
}
static int lib_decompress_create(lua_State *ctx, int window_size) {
	// size_t data_size;
	// uint8_t *data = (uint8_t*)luaL_checklstring(ctx, 1, &data_size);

	lua_pushvalue(ctx, 1);

	lua_zlib_stream_t *stream = lua_newuserdatauv(ctx, sizeof *stream, 0);
	memset(stream, 0, sizeof *stream);

	luaL_getmetatable(ctx, DECOMPRESS_META);
	lua_setmetatable(ctx, -2);


	// stream->stream.avail_out = sizeof stream->buff;
	// stream->stream.next_out = stream->buff;

	// stream->stream.avail_in = data_size;
	// stream->stream.next_in = data;

	if (inflateInit2(&stream->stream, window_size) != Z_OK) {
		return luaL_error(ctx, "failed to initialize inflate stream");
	}

	lua_pushcclosure(ctx, lib_decompress_iter, 2);
	return 1;
}


static int lib_inflate(lua_State *ctx) {
	return lib_decompress_create(ctx, MAX_WBITS);
}
static int lib_inflate_raw(lua_State *ctx) {
	return lib_decompress_create(ctx, -MAX_WBITS);
}
static int lib_gunzip(lua_State *ctx) {
	return lib_decompress_create(ctx, MAX_WBITS + 16);
}
static int lib_decompress(lua_State *ctx) {
	return lib_decompress_create(ctx, MAX_WBITS + 32);
}


int zlib_open_lib(lua_State *ctx) {
	luaL_newmetatable(ctx, DECOMPRESS_META);
	luaL_setfuncs(ctx, (luaL_Reg[]) {
		{ "__gc", lib_decompress_free },
		{ NULL, NULL },
	}, 0);

	lua_pop(ctx, 2);

	luaL_newlib(ctx, ((luaL_Reg[]) {
		{ "inflate", lib_inflate },
		{ "inflate_raw", lib_inflate_raw },
		{ "gunzip", lib_gunzip },
		{ "decompress", lib_decompress },
		{ NULL, NULL },
	}));

	return 1;
}