#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <time.h>

/*
** match allocator’s thresholds
** (only used here for shaping size distributions)
*/
#define TINY_MAX 64
#define SMALL_MAX 4096

#define ARRAY_SIZE 10000	// number of live slots we manage
#define RANDOM_ITERS 200000 // number of random operations
#define PATTERN_BASE 0xA5
#define MAX_LARGE_SIZE (SMALL_MAX * 8) // up to 32k-ish

typedef struct s_block
{
	void *ptr;
	size_t size;
	unsigned char pattern;
	int is_calloc;
} t_block;

static t_block g_blocks[ARRAY_SIZE];

/* Fill a block with a simple pattern byte */
static void fill_pattern(void *ptr, size_t size, unsigned char pattern)
{
	if (!ptr || size == 0)
		return;
	memset(ptr, pattern, size);
}

/* Check that a block still has the expected pattern */
static int check_pattern(void *ptr, size_t size, unsigned char pattern)
{
	if (!ptr || size == 0)
		return 1;
	unsigned char *p = (unsigned char *)ptr;
	for (size_t i = 0; i < size; ++i)
	{
		if (p[i] != pattern)
		{
			fprintf(stderr,
					"[ERROR] Pattern mismatch at %p (offset %zu): "
					"expected 0x%02X, got 0x%02X\n",
					ptr, i, (unsigned)pattern, (unsigned)p[i]);
			return 0;
		}
	}
	return 1;
}

/* Check that a calloc'd block is zeroed */
static int check_zero(void *ptr, size_t size)
{
	if (!ptr || size == 0)
		return 1;
	unsigned char *p = (unsigned char *)ptr;
	for (size_t i = 0; i < size; ++i)
	{
		if (p[i] != 0)
		{
			fprintf(stderr,
					"[ERROR] Calloc block not zeroed at %p (offset %zu): got 0x%02X\n",
					ptr, i, (unsigned)p[i]);
			return 0;
		}
	}
	return 1;
}

/* Generate a pseudo-random size skewed across tiny/small/large */
static size_t random_size(void)
{
	int r = rand() % 100;

	if (r < 40)
	{
		/* TINY range [1..TINY_MAX] */
		return (size_t)(1 + rand() % TINY_MAX);
	}
	else if (r < 85)
	{
		/* SMALL range [TINY_MAX+1 .. SMALL_MAX] */
		return (size_t)(TINY_MAX + 1 + rand() % (SMALL_MAX - TINY_MAX));
	}
	else
	{
		/* LARGE range [SMALL_MAX+1 .. MAX_LARGE_SIZE] */
		return (size_t)(SMALL_MAX + 1 + rand() % (MAX_LARGE_SIZE - SMALL_MAX - 1));
	}
}

/* Tear down all blocks cleanly (for the end of tests) */
static void free_all_blocks(void)
{
	for (int i = 0; i < ARRAY_SIZE; ++i)
	{
		if (g_blocks[i].ptr)
		{
			free(g_blocks[i].ptr);
			g_blocks[i].ptr = NULL;
			g_blocks[i].size = 0;
		}
	}
}

/* Simple deterministic sanity tests */
static void basic_tests(void)
{
	printf("== basic_tests ==\n");

	/* TINY */
	void *a = malloc(10);
	void *b = malloc(TINY_MAX);
	fill_pattern(a, 10, 0x11);
	fill_pattern(b, TINY_MAX, 0x22);
	check_pattern(a, 10, 0x11);
	check_pattern(b, TINY_MAX, 0x22);
	free(a);
	free(b);

	/* SMALL */
	void *c = malloc(SMALL_MAX / 2);
	fill_pattern(c, SMALL_MAX / 2, 0x33);
	check_pattern(c, SMALL_MAX / 2, 0x33);
	c = realloc(c, SMALL_MAX / 2 + 100); // grow within SMALL
	check_pattern(c, SMALL_MAX / 2, 0x33);
	free(c);

	/* LARGE */
	size_t large_sz = SMALL_MAX * 4;
	void *d = malloc(large_sz);
	fill_pattern(d, large_sz, 0x44);
	check_pattern(d, large_sz, 0x44);
	d = realloc(d, large_sz * 2); // grow LARGE
	check_pattern(d, large_sz, 0x44);
	d = realloc(d, large_sz / 2); // shrink LARGE
	check_pattern(d, large_sz / 2, 0x44);
	free(d);

	/* CALLOC */
	void *e = calloc(1000, 1);
	check_zero(e, 1000);
	memset(e, 0xAB, 1000);
	e = realloc(e, 2000);
	check_pattern(e, 1000, 0xAB);
	free(e);

	printf("basic_tests: OK\n\n");
}

static void random_stress(void)
{
	printf("== random_stress ==\n");

	for (int i = 0; i < ARRAY_SIZE; ++i)
	{
		g_blocks[i].ptr = NULL;
		g_blocks[i].size = 0;
		g_blocks[i].pattern = 0;
		g_blocks[i].is_calloc = 0;
	}

	for (int iter = 0; iter < RANDOM_ITERS; ++iter)
	{
		int idx = rand() % ARRAY_SIZE;
		t_block *blk = &g_blocks[idx];
		int op = rand() % 4; /* 0 = alloc/new, 1 = free, 2 = realloc up, 3 = realloc down */

		/* Invariant: for non-NULL blocks we always keep full [0..size) filled with pattern */
		if (blk->ptr && !blk->is_calloc)
		{
			if (!check_pattern(blk->ptr, blk->size, blk->pattern))
			{
				fprintf(stderr, "[FATAL] pattern mismatch before op; idx=%d\n", idx);
				abort();
			}
		}

		if (op == 0)
		{
			/* allocate new if empty, otherwise grow via realloc */
			if (blk->ptr == NULL)
			{
				size_t sz = random_size();
				int use_calloc = (rand() % 4) == 0; /* 25% calloc */

				if (use_calloc)
				{
					blk->ptr = calloc(sz, 1);
					blk->is_calloc = 1;
					blk->size = sz;
					if (!blk->ptr)
					{
						fprintf(stderr, "[ERROR] calloc returned NULL\n");
						continue;
					}
					/* we only check zero once; then we switch to pattern mode */
					if (!check_zero(blk->ptr, blk->size))
						fprintf(stderr, "[ERROR] calloc block not zeroed (idx=%d)\n", idx);
					blk->pattern = (unsigned char)((PATTERN_BASE + idx) & 0xFF);
					fill_pattern(blk->ptr, blk->size, blk->pattern);
					blk->is_calloc = 0;
				}
				else
				{
					blk->ptr = malloc(sz);
					blk->is_calloc = 0;
					blk->size = sz;
					if (!blk->ptr)
					{
						fprintf(stderr, "[ERROR] malloc returned NULL\n");
						continue;
					}
					blk->pattern = (unsigned char)((PATTERN_BASE + idx) & 0xFF);
					fill_pattern(blk->ptr, blk->size, blk->pattern);
				}
			}
			else
			{
				/* grow via realloc */
				size_t old_size = blk->size;
				size_t new_sz = old_size + 1 + (rand() % (old_size + 1));
				unsigned char pattern = blk->pattern;

				void *new_ptr = realloc(blk->ptr, new_sz);
				if (!new_ptr)
				{
					fprintf(stderr, "[ERROR] realloc (grow) returned NULL\n");
					continue; /* keep old block as-is */
				}
				blk->ptr = new_ptr;
				blk->size = new_sz;
				blk->pattern = pattern;

				/* Only the first old_size bytes must be preserved */
				if (!check_pattern(blk->ptr, old_size, pattern))
				{
					fprintf(stderr, "[ERROR] pattern mismatch after realloc grow (idx=%d)\n", idx);
					abort();
				}
				/* Now enforce our invariant for the future */
				fill_pattern(blk->ptr, blk->size, blk->pattern);
			}
		}
		else if (op == 1)
		{
			/* free */
			if (blk->ptr)
			{
				free(blk->ptr);
				blk->ptr = NULL;
				blk->size = 0;
				blk->is_calloc = 0;
			}
		}
		else if (op == 2)
		{
			/* explicit realloc up (if exists) */
			if (blk->ptr)
			{
				size_t old_size = blk->size;
				size_t new_sz = old_size + 1 + (rand() % (old_size + 1));
				unsigned char pattern = blk->pattern;

				void *new_ptr = realloc(blk->ptr, new_sz);
				if (!new_ptr)
				{
					fprintf(stderr, "[ERROR] realloc (grow2) returned NULL\n");
					continue;
				}

				blk->ptr = new_ptr;
				blk->size = new_sz;
				blk->pattern = pattern;

				if (!check_pattern(blk->ptr, old_size, pattern))
				{
					fprintf(stderr, "[ERROR] pattern mismatch after realloc grow2 (idx=%d)\n", idx);
					abort();
				}
				fill_pattern(blk->ptr, blk->size, blk->pattern);
			}
		}
		else /* op == 3, realloc down */
		{
			if (blk->ptr && blk->size > 1)
			{
				size_t old_size = blk->size;
				size_t new_sz = 1 + (rand() % blk->size);
				unsigned char pattern = blk->pattern;

				void *new_ptr = realloc(blk->ptr, new_sz);
				if (!new_ptr)
				{
					fprintf(stderr, "[ERROR] realloc (shrink) returned NULL\n");
					continue;
				}

				blk->ptr = new_ptr;
				blk->size = new_sz;
				blk->pattern = pattern;

				/* For shrink, we only expect the first new_sz bytes to still match */
				if (!check_pattern(blk->ptr, blk->size, blk->pattern))
				{
					fprintf(stderr, "[ERROR] pattern mismatch after realloc shrink (idx=%d)\n", idx);
					abort();
				}
				/* Still keep invariant explicitly (optional here) */
				fill_pattern(blk->ptr, blk->size, blk->pattern);
			}
		}

		if ((iter % 50000) == 0 && iter != 0)
			printf("  random_stress progress: %d / %d\n", iter, RANDOM_ITERS);
	}

	free_all_blocks();
	printf("random_stress: DONE\n\n");
}

/* Fragmentation test: allocate many small then free some, then big allocs, etc. */
static void fragmentation_test(void)
{
	printf("== fragmentation_test ==\n");

	const int count = 10000;
	void **arr = malloc(sizeof(void *) * count);
	size_t *sizes = malloc(sizeof(size_t) * count);
	if (!arr || !sizes)
	{
		fprintf(stderr, "[ERROR] fragmentation_test: malloc failed\n");
		free(arr);
		free(sizes);
		return;
	}

	/* Step 1: allocate random small blocks */
	for (int i = 0; i < count; ++i)
	{
		size_t sz = 1 + (rand() % SMALL_MAX);
		arr[i] = malloc(sz);
		sizes[i] = sz;
		if (!arr[i])
		{
			fprintf(stderr, "[ERROR] fragmentation_test: malloc failed at i=%d\n", i);
			continue;
		}
		memset(arr[i], 0x5A, sz);
	}

	/* Step 2: free every other block */
	for (int i = 0; i < count; i += 2)
	{
		if (arr[i])
		{
			free(arr[i]);
			arr[i] = NULL;
			sizes[i] = 0;
		}
	}

	/* Step 3: allocate a bunch of bigger blocks, try to reuse fragmentation */
	const int big_count = 1000;
	void **big = malloc(sizeof(void *) * big_count);
	size_t *big_sz = malloc(sizeof(size_t) * big_count);
	if (!big || !big_sz)
	{
		fprintf(stderr, "[ERROR] fragmentation_test: big malloc failed\n");
		free(big);
		free(big_sz);
		free(arr);
		free(sizes);
		return;
	}

	for (int i = 0; i < big_count; ++i)
	{
		size_t sz = SMALL_MAX + (rand() % (SMALL_MAX * 4));
		big[i] = malloc(sz);
		big_sz[i] = sz;
		if (!big[i])
		{
			fprintf(stderr, "[ERROR] fragmentation_test: big malloc failed at i=%d\n", i);
			continue;
		}
		memset(big[i], 0x6B, sz);
	}

	/* cleanup */
	for (int i = 0; i < count; ++i)
		free(arr[i]);
	for (int i = 0; i < big_count; ++i)
		free(big[i]);
	free(arr);
	free(sizes);
	free(big);
	free(big_sz);

	printf("fragmentation_test: DONE\n\n");
}

int main(void)
{
	srand((unsigned int)time(NULL));

	basic_tests();
	random_stress();
	fragmentation_test();

	printf("ALL TESTS DONE\n");
	return 0;
}