#include "catch.hpp"
#include "test_helpers.hpp"
#include "duckdb/main/appender.hpp"
#include "duckdb/common/atomic.hpp"

using namespace duckdb;
using namespace std;

struct MyAllocateData : public PrivateAllocatorData {
	MyAllocateData(atomic<idx_t> *memory_counter_p) : memory_counter(memory_counter_p) {
	}

	atomic<idx_t> *memory_counter;
};

data_ptr_t my_allocate_function(PrivateAllocatorData *private_data, idx_t size) {
	auto my_allocate_data = (MyAllocateData *)private_data;
	*my_allocate_data->memory_counter += size;
	return data_ptr_cast(malloc(size));
}

void my_free_function(PrivateAllocatorData *private_data, data_ptr_t pointer, idx_t size) {
	auto my_allocate_data = (MyAllocateData *)private_data;
	*my_allocate_data->memory_counter -= size;
	free(pointer);
}

data_ptr_t my_reallocate_function(PrivateAllocatorData *private_data, data_ptr_t pointer, idx_t old_size, idx_t size) {
	auto my_allocate_data = (MyAllocateData *)private_data;
	*my_allocate_data->memory_counter -= old_size;
	*my_allocate_data->memory_counter += size;
	return data_ptr_cast(realloc(pointer, size));
}

TEST_CASE("Test using a custom allocator", "[api][.]") {
	atomic<idx_t> memory_counter;
	memory_counter = 0;

	REQUIRE(memory_counter.load() == 0);

	DBConfig config;
	config.allocator = make_uniq<Allocator>(my_allocate_function, my_free_function, my_reallocate_function,
	                                        make_uniq<MyAllocateData>(&memory_counter));
	DuckDB db(nullptr, &config);
	Connection con(db);
	REQUIRE_NO_FAIL(con.Query("CREATE TABLE tbl AS SELECT * FROM range(1000000)"));

	// check that the memory counter reported anything
	REQUIRE(memory_counter.load() > 0);
	auto table_memory_usage = memory_counter.load();

	REQUIRE_NO_FAIL(con.Query("DROP TABLE tbl"));

	// check that the memory counter usage has decreased after we dropped the table
	REQUIRE(memory_counter.load() < table_memory_usage);
}