#define _DEFAULT_SOURCE
#include <assert.h>
#include <errno.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define ALIGN(x, a) (((x) + (a)-1) & ~(((typeof(x))(a)-1)))
#define MIN_ALLOC_SIZE (1 << 12) // 4 kbytes.

#define PTR_TO_CHUNK(ptr) ((free_chunk*)(((size_t*)ptr) - 1))
#define CHUNK_TO_PTR(chunk) (((size_t*)chunk) + 1)

#ifndef NDEBUG
#define debug fprintf
#else
#define debug(...) (void)0;
#endif

//#define NEXT_CHUNK(ch) ((chunk*)((uint8_t*)((ch) + 1) + GET_SIZE(ch)))
#define CHUNK_PLUS(ch, sz) ((free_chunk*)((uint8_t*)(((size_t*)ch) + 1) + (sz)))

enum flags { FLAG_UNUSED = 1, FLAG_INUSE = 2 };

typedef struct free_chunk {
    size_t size;
    /* members after size are in the 'data' area of an unused chunk. */
    struct free_chunk* next;
} free_chunk;

/* A pair of chunks. Just used as a return value for easy unlinking. */
typedef struct chunk_pair {
    free_chunk* predecessor;
    free_chunk* chunk;
} chunk_pair;

static free_chunk* free_chunk_head = NULL;
static free_chunk* free_chunk_tail = NULL;

/* Prints the freelist, for debugging purposes. */
void print_freelist() {
    if (!free_chunk_head) {
        assert(free_chunk_tail == NULL);
        debug(stderr, "Freelist is empty.\n");
        return;
    }

    int i = 1;
    for (free_chunk* fc = free_chunk_head; fc; fc = fc->next, i++) {
        debug(stderr, "Chunk %d at %p: size %lu, next at %p\n", i, fc, fc->size,
              fc->next);
    }

    debug(stderr, "Chunk tail is %p (should be same as last chunk).\n",
          free_chunk_tail);
}


/* Finds a predecessor for a chunk.
 * You probably don't need to use this directly. */
static free_chunk* _find_predecessor(free_chunk* chunk) {
    if (chunk < free_chunk_head)
        return NULL;

    for (free_chunk* ch = free_chunk_head; ch; ch = ch->next) {
        if (ch < chunk && ch->next > chunk) {
            return ch;
        }
    }

    /* Well, that was just an expensive way of finding the tail of the list.. */
    print_freelist();
    assert(free_chunk_tail != NULL);
    assert(chunk > free_chunk_tail);
    return free_chunk_tail;
}

/* Insert a given chunk into the 'right' place on the freelist (ordered by
 * address). Potentially updates the head or tail of the freelist.
 *
 * If hint != NULL, inserts the chunk AFTER hint.
 */
static void insert_free_chunk(free_chunk* to_insert, free_chunk* hint) {
    if (!free_chunk_head) {
        free_chunk_head = free_chunk_tail = to_insert;
        to_insert->next = NULL;
    } else if (to_insert > free_chunk_tail) {
        free_chunk_tail->next = to_insert;
        to_insert->next = NULL;
        free_chunk_tail = to_insert;
    } else if (to_insert < free_chunk_head) {
        to_insert->next = free_chunk_head;
        free_chunk_head = to_insert;
    } else if (hint) {
        to_insert->next = hint->next;
        hint->next = to_insert;
    } else {
        free_chunk* predecessor = _find_predecessor(to_insert);
        to_insert->next = predecessor->next;
        predecessor->next = to_insert;
    }
}

/* Unlinks a free chunk from the free list. */
static void unlink_free_chunk(free_chunk* previous, free_chunk* current) {
    /* Since this is singly linked, unlinking is slightly more difficult,
     * because we need a previous pointer (or need to traverse the entire list)
     * It's worth it because of the space constraints, though. */
    if (current == free_chunk_head)
        free_chunk_head = current->next;
    else
        previous->next = current->next;

    if (current == free_chunk_tail)
        free_chunk_tail = previous;

    current->next = NULL;
}

/* Allocates a new chunk, of min(MIN_ALLOC_SIZE, size) bytes.
 * Returns the chunk if successful, NULL otherwise.
 */
static free_chunk* allocate_new_chunk(size_t size) {
#ifdef DEBUG_ALLOCATE_AMOUNT
    debug(stderr,
          "Allocating new chunk with at least %lu bytes (without header)\n",
          size > MIN_ALLOC_SIZE ? size : MIN_ALLOC_SIZE);
#endif

    size_t alloc_size = size;
    if (size < MIN_ALLOC_SIZE)
        alloc_size = MIN_ALLOC_SIZE;

    void* base = sbrk(alloc_size + sizeof(size_t));

    if (base == (void*)-1)
        return NULL;

    free_chunk* ch = base;
    ch->size = alloc_size;

    return ch;
}

static chunk_pair find_chunk_with_space(size_t space) {
    /* Try to find something in the freelist. */
    free_chunk* previous = NULL;
    for (free_chunk* fc = free_chunk_head; fc; previous = fc, fc = fc->next) {
        if (fc->size >= space) {
            /* Take ch out of the freelist. */
            unlink_free_chunk(previous, fc);

            return (chunk_pair){previous, fc};
        }
    }

    /* If nothing was found, allocate a new chunk. */

    free_chunk* allocated;
    if (!(allocated = allocate_new_chunk(space)))
        /* Well, that failed. Return NULL. */
        return (chunk_pair){NULL, NULL};

    return (chunk_pair){NULL, allocated};
}

void* mymalloc(size_t size) {
    if (!size) {
        /* Grep depends on this returning non-null. */
        return (void*)0x1;
    }

    size_t aligned_size = ALIGN(size, sizeof(long));

    chunk_pair cpair = find_chunk_with_space(aligned_size);
    free_chunk* ch = cpair.chunk;
    free_chunk* pred = cpair.predecessor;

    if (!ch) {
        debug(stderr, "No block with space found or creatable!\n");
        errno = ENOMEM;
        return NULL;
    }

    size_t chunk_size = ch->size;

    if (chunk_size > aligned_size + sizeof(free_chunk)) {
        // There is enough space to split this chunk into two.
        free_chunk* new_chunk = CHUNK_PLUS(ch, aligned_size);
        new_chunk->size = chunk_size - aligned_size - sizeof(size_t);

        /* This shouldn't overflow. */
        assert(new_chunk->size < chunk_size);

        /* Insert the free chunk where ch used to be. */
        insert_free_chunk(new_chunk, pred);

        ch->size = aligned_size;
    }

    void *ptr = CHUNK_TO_PTR(ch);

#ifdef DEBUG_CALL_RESULTS
    debug(stderr, "Returning pointer: %p from chunk %p\n", ptr, ch);
#endif

    errno = 0;

    return ptr;
}

void* mycalloc(size_t nmemb, size_t size) {
    if (nmemb == 0 || size == 0) {
        errno = ENOMEM;
        return NULL;
    }

    void* ptr;
    size_t sz = nmemb * size;

    if (sz / size != nmemb) {
        // Overflow
        errno = ENOMEM;
        return NULL;
    }
    ptr = mymalloc(sz);
    if (ptr)
        memset(ptr, 0, sz);
    errno = 0;
    return ptr;
}

static inline void print_free_chunk_info(free_chunk* ch) {
    debug(stderr, "Free chunk: size %lu, next at %p\n", ch->size, ch->next);
}

static inline void print_chunk_info(size_t* ch) {
    debug(stderr, "Chunk: size %lu\n", *ch);
}

/*
static void merge_chunks(free_chunk* a, free_chunk* b) {
    // We'll assume that a, b are in that order in memory.
    a->size += b->size + sizeof(size_t);
    unlink_free_chunk(a, b);
}

static void try_merge_adjacent() {
    return;
    chunk* previous = NULL;
    int i = 1;
    CHUNK_FOR(ch) {
        // print_chunk_info(ch);
        if (previous && GET_FLAGS(previous) == GET_FLAGS(ch) &&
            GET_FLAGS(ch) == FLAG_UNUSED) {
            debug(stderr, "Merging chunks at %p and %p (sizes %lu, %lu)\n",
                  previous, ch, GET_SIZE(previous), GET_SIZE(ch));
            ch = merge_chunks(previous, ch);
        }
        previous = ch;
        i++;
    }
} */

void myfree(void* ptr) {
    /* Grep, again. And just a normal NULL check. */
    if (!ptr || ptr == (void*)0x1)
        return;

#ifdef DEBUG_CALL_RESULTS
    debug(stderr, "Free called with ptr %p\n", ptr);
#endif
    free_chunk* chunk = PTR_TO_CHUNK(ptr);

    /* Insert the chunk into the free list. */
    free_chunk* hint = _find_predecessor(chunk);
    insert_free_chunk(chunk, hint);

    // try_merge_adjacent();
    errno = 0;
}

void* myrealloc(void* ptr, size_t size) {
    debug(stderr, "Realloc called with %p, %lu\n", ptr, size);

    if (ptr != NULL && size == 0) {
        // If size is equal to zero, and ptr != NULL, then the call is
        // equivalent to free(ptr)
        debug(stderr, "realloc: returning NULL as size == 0\n");
        myfree(ptr);
        return NULL;
    }

    if (!ptr) {
        debug(stderr, "realloc: returning new block with size %lu\n", size);
        return mymalloc(size);
    }

    free_chunk* ch = PTR_TO_CHUNK(ptr);
    size_t alloc_size = ch->size;

    debug(stderr, "Allocation size: %lu\n", alloc_size);

    if (alloc_size >= size) {
        debug(stderr,
              "realloc: current block is large enough. returning that.\n");
        errno = 0;
        return ptr;
    }

    void* newptr;
    newptr = mymalloc(size);
    if (newptr)
        memcpy(newptr, ptr, alloc_size);

    myfree(ptr);

    errno = 0;
    return newptr;
}

void* malloc(size_t size) {
    void* ptr = mymalloc(size);
    return ptr;
}
void* calloc(size_t nmemb, size_t size) {
    void* ptr = mycalloc(nmemb, size);
    return ptr;
}
void* realloc(void* ptr, size_t size) {
    void* newptr = myrealloc(ptr, size);
    return newptr;
}
void free(void* ptr) {
    myfree(ptr);
}