/* radare - LGPL - Copyright 2022-2024 - condret */

#include <r_io.h>

typedef struct io_treebuf_t {
	RRBTree *tree;
	ut64 seek;
} IOTreeBuf;

typedef struct io_treebuf_chunk_t {
	RInterval itv;
	ut8 *buf;
} IOTreeBufChunk;

static void _treebuf_chunk_free(void *data) {
	IOTreeBufChunk *chunk = (IOTreeBufChunk *)data;
	if (chunk) {
		free (chunk->buf);
		free (chunk);
	}
}

static bool __check(RIO *io, const char *pathname, bool many) {
	return !strcmp (pathname, "treebuf://");
}

static char *__system(RIO *io, RIODesc *desc, const char *cmd) {
	if (cmd && !strcmp (cmd, "reset")) {
		IOTreeBuf *treebuf = (IOTreeBuf *)desc->data;
		r_crbtree_free (treebuf->tree);
		treebuf->tree = r_crbtree_new (_treebuf_chunk_free);
	}
	return NULL;
}

static RIODesc *__open(RIO *io, const char *pathname, int rw, int mode) {
	if (!__check (io, pathname, false)) {
		return NULL;
	}
	IOTreeBuf *treebuf = R_NEW0 (IOTreeBuf);
	treebuf->tree = r_crbtree_new (_treebuf_chunk_free);
	RIODesc *desc = r_io_desc_new (io, &r_io_plugin_treebuf, pathname, R_PERM_RW, mode, treebuf);
	if (!desc) {
		r_crbtree_free (treebuf->tree);
		free (treebuf);
	}
	return desc;
}

static bool __close(RIODesc *desc) {
	IOTreeBuf *treebuf = desc->data;
	r_crbtree_free (treebuf->tree);
	R_FREE (desc->data);
	return true;
}

static ut64 __lseek(RIO* io, RIODesc *desc, ut64 offset, int whence) {
	IOTreeBuf *treebuf = (IOTreeBuf *)desc->data;
	switch (whence) {
	case R_IO_SEEK_SET:
		return treebuf->seek = offset;
	case R_IO_SEEK_END:
		return treebuf->seek = UT64_MAX - 1;
	case R_IO_SEEK_CUR:
		return treebuf->seek = R_MAX (treebuf->seek, treebuf->seek + offset);
	}
	R_LOG_ERROR ("Invalid whence %d", whence);
	return treebuf->seek;
}

static int _treebuf_chunk_find(void *incoming, void *in, void *user) {
	RInterval *itv = (RInterval *)incoming;
	IOTreeBufChunk *chunk = (IOTreeBufChunk *)in;
	if (r_itv_overlap (itv[0], chunk->itv)) {
		return 0;
	}
	return (r_itv_begin (itv[0]) < r_itv_begin (chunk->itv))? -1: 1;
}

static int _treebuf_chunk_insert(void *incoming, void *in, void *user) {
	const ut64 ia = r_itv_begin (((IOTreeBufChunk *)incoming)->itv);
	const ut64 na = r_itv_begin (((IOTreeBufChunk *)in)->itv);
	if (ia < na) {
		return -1;
	}
	return (ia > na)? 1: 0;
}

static RRBNode *_first_overlap(RRBNode *node, RInterval itv) {
	RRBNode *prev = r_rbnode_prev (node);
	while (prev) {
		IOTreeBufChunk *chunk = (IOTreeBufChunk *)prev->data;
		if (!r_itv_overlap (chunk->itv, itv)) {
			break;
		}
		node = prev;
		prev = r_rbnode_prev (prev);
	}
	return node;
}

static int _treebuf_new_chunk(IOTreeBuf *treebuf, RInterval itv, const ut8 *buf) {
	IOTreeBufChunk *chunk = R_NEW0 (IOTreeBufChunk);
	chunk->itv = itv;
	chunk->buf = R_NEWS (ut8, itv.size);
	if (!chunk->buf || !r_crbtree_insert (treebuf->tree, chunk, _treebuf_chunk_insert, NULL)) {
		free (chunk->buf);
		free (chunk);
		return -1;
	}
	memcpy (chunk->buf, buf, itv.size);
	treebuf->seek = r_itv_end (itv);
	return (int)itv.size;
}

static int __read(RIO *io, RIODesc *desc, ut8 *buf, int len) {
	IOTreeBuf *treebuf = (IOTreeBuf *)desc->data;
	RInterval search_itv = {treebuf->seek, R_MIN ((ut64)len, UT64_MAX - treebuf->seek)};
	memset (buf, io->Oxff, len);
	RRBNode *node = r_crbtree_find_node (treebuf->tree, &search_itv, _treebuf_chunk_find, NULL);
	if (!node) {
		treebuf->seek = r_itv_end (search_itv);
		return (int)r_itv_size (search_itv);
	}
	node = _first_overlap (node, search_itv);
	IOTreeBufChunk *chunk = (IOTreeBufChunk *)node->data;
	const ut64 sbeg = r_itv_begin (search_itv);
	const ut64 send = r_itv_end (search_itv);
	do {
		const ut64 cbeg = r_itv_begin (chunk->itv);
		const ut64 addr = R_MAX (sbeg, cbeg);
		const ut64 end = R_MIN (send, r_itv_end (chunk->itv));
		memcpy (&buf[addr - sbeg], &chunk->buf[addr - cbeg], (size_t)(end - addr));
		node = r_rbnode_next (node);
		chunk = node? (IOTreeBufChunk *)node->data: NULL;
	} while (chunk && r_itv_overlap (chunk->itv, search_itv));
	treebuf->seek = send;
	return (int)r_itv_size (search_itv);
}

static int __write(RIO *io, RIODesc *desc, const ut8 *buf, int len) {
	IOTreeBuf *treebuf = (IOTreeBuf *)desc->data;
	RInterval search_itv = {treebuf->seek, R_MIN ((ut64)len, UT64_MAX - treebuf->seek)};
	RRBNode *node = r_crbtree_find_node (treebuf->tree, &search_itv, _treebuf_chunk_find, NULL);
	if (!node) {
		return _treebuf_new_chunk (treebuf, search_itv, buf);
	}
	node = _first_overlap (node, search_itv);
	IOTreeBufChunk *chunk = (IOTreeBufChunk *)node->data;
	if (r_itv_include (chunk->itv, search_itv)) {
		memcpy (&chunk->buf[r_itv_begin (search_itv) - r_itv_begin (chunk->itv)],
			buf, r_itv_size (search_itv));
		treebuf->seek = r_itv_end (search_itv);
		return (int)r_itv_size (search_itv);
	}
	if (r_itv_begin (chunk->itv) < r_itv_begin (search_itv)) {
		chunk->itv.size = r_itv_begin (search_itv) - r_itv_begin (chunk->itv);
		chunk->buf = realloc (chunk->buf, r_itv_size (chunk->itv));
		node = r_rbnode_next (node);
	}
	if (node) {
		chunk = (IOTreeBufChunk *)node->data;
		while (chunk && r_itv_include (search_itv, chunk->itv)) {
			node = r_rbnode_next (node);
			r_crbtree_delete (treebuf->tree, &chunk->itv, _treebuf_chunk_find, NULL);
			chunk = node? (IOTreeBufChunk *)node->data: NULL;
		}
		if (chunk && r_itv_end (search_itv) >= r_itv_begin (chunk->itv)) {
			chunk->buf = realloc (chunk->buf, r_itv_end (chunk->itv) - r_itv_begin (search_itv));
			memmove (&chunk->buf[r_itv_size (search_itv)],
				&chunk->buf[r_itv_end (search_itv) - r_itv_begin (chunk->itv)],
				r_itv_end (chunk->itv) - r_itv_end (search_itv));
			memcpy (chunk->buf, buf, r_itv_size (search_itv));
			chunk->itv.size = r_itv_end (chunk->itv) - r_itv_begin (search_itv);
			chunk->itv.addr = search_itv.addr;
			treebuf->seek = r_itv_end (search_itv);
			return (int)r_itv_size (search_itv);
		}
	}
	return _treebuf_new_chunk (treebuf, search_itv, buf);
}

RIOPlugin r_io_plugin_treebuf = {
	.meta = {
		.name = "treebuf",
		.author = "condret",
		.desc = "Dynamic sparse like buffer without size restriction",
		.license = "LGPL-3.0-only",
	},
	.uris = "treebuf://",
	.system = __system,
	.open = __open,
	.close = __close,
	.read = __read,
	.check = __check,
	.seek = __lseek,
	.write = __write,
};

#ifndef R2_PLUGIN_INCORE
R_API RLibStruct radare_plugin = {
	.type = R_LIB_TYPE_IO,
	.data = &r_io_plugin_treebuf,
	.version = R2_VERSION
};
#endif