#ifndef STACK_H
#define STACK_H

#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#if defined(__GNUC__) || defined(__clang__)
#define MUST_CHECK __attribute__((warn_unused_result))
#else
#define MUST_CHECK
#endif

typedef struct {
    uint8_t *buf;     //!< Pointer to the internal buffer
    size_t   buf_sz;  //!< The length of the allocated buffer, in bytes
    size_t   idx;     //!< Indicates where to read/write
    size_t   item_sz; //!< The size of the contained item, in bytes
} Stack;

typedef enum {
    Ok = 0,         //!< Everything is fine
    Error = -1,     //!< Something went wrong
    StackFull = 1,  //!< The stack is full
    StackEmpty = 2, //!< The stack is empty
} StackResult;

MUST_CHECK StackResult stack_init(Stack *stack, size_t capacity, size_t item_sz);
MUST_CHECK StackResult stack_init_raw(Stack *stack, uint8_t *buf, size_t buf_sz, size_t item_sz);
MUST_CHECK StackResult stack_resize(Stack *stack, size_t capacity);
uint8_t               *stack_resize_raw(Stack *stack, uint8_t *new_buf, size_t buf_sz);
MUST_CHECK StackResult stack_destroy(Stack *stack);
MUST_CHECK StackResult stack_push(Stack *stack, const void *value);
MUST_CHECK StackResult stack_pop(Stack *stack, void *dest);

static inline size_t stack_length(Stack *stack) {
    return (stack->idx / stack->item_sz);
}

static inline size_t stack_capacity(Stack *stack) {
    return (stack->buf_sz / stack->item_sz);
}

static inline bool stack_full(Stack *stack) {
    return (stack->buf_sz == stack->idx);
}

static inline bool stack_empty(Stack *stack) {
    return (stack->idx < stack->item_sz);
}

#endif // STACK_H

// #include "stack.h"

StackResult stack_init(Stack *stack, size_t capacity, size_t item_sz) {
    return stack_init_raw(stack, malloc(capacity * item_sz), capacity * item_sz, item_sz);
}

StackResult stack_init_raw(Stack *stack, uint8_t *buf, size_t buf_sz, size_t item_sz) {
    if (!memset(stack, 0, sizeof(Stack))) {
        return Error;
    }

    stack->buf = buf;
    if (!stack->buf) {
        return Error;
    }

    stack->buf_sz = buf_sz;
    stack->item_sz = item_sz;
    stack->idx = 0;

    memset(stack->buf, 0, buf_sz);

    return Ok;
}

StackResult stack_resize(Stack *stack, size_t capacity) {
    uint8_t *old_buf = stack_resize_raw(stack, malloc(capacity * stack->item_sz), capacity * stack->item_sz);

    if (old_buf)
        free(old_buf);

    return Ok;
}

uint8_t *stack_resize_raw(Stack *stack, uint8_t *new_buf, size_t buf_sz) {
    if (stack->idx > buf_sz)
        return NULL;

    uint8_t *old_buf = stack->buf;
    memcpy(new_buf, old_buf, stack->buf_sz);

    stack->buf = new_buf;
    stack->buf_sz = buf_sz;

    return old_buf;
}

StackResult stack_push(Stack *stack, const void *value) {
    if (stack->idx + stack->item_sz > stack->buf_sz)
        return StackFull;

    memcpy((char *)stack->buf + stack->idx, value, stack->item_sz);
    stack->idx += stack->item_sz;

    return Ok;
}

StackResult stack_pop(Stack *stack, void *dest) {
    if (stack->idx < stack->item_sz)
        return StackEmpty;

    stack->idx -= stack->item_sz;
    memcpy(dest, (char *)stack->buf + stack->idx, stack->item_sz);

    return Ok;
}

StackResult stack_destroy(Stack *stack) {
    free(stack->buf);

    stack->buf = NULL;
    stack->buf_sz = 0;
    stack->item_sz = 0;
    stack->idx = 0;

    return Ok;
}

int main(void) {
    Stack s;
    if (stack_init(&s, 128, sizeof(int)) != Ok)
        return 1;

    assert(stack_capacity(&s) == 128);

    int r; // Just used for testing

    { /* Sanity check: push some data and check the output */
        for (int i = 1; i <= 10; i++) {
            if (stack_push(&s, &i) != Ok)
                return 1;
        }

        while (stack_pop(&s, &r) == Ok) {
            printf("Number %d\n", r);
        }
    }

    /* Should tell us StackEmpty */
    assert(stack_pop(&s, &r) == StackEmpty);

    /* Make sure it tolerates filling and behaves expectedly */
    while (stack_push(&s, &r) == Ok);
    assert(stack_full(&s));
    assert(stack_length(&s) == 128);
    assert(stack_push(&s, &r) == StackFull);

    /* Drain it and check for proper behaviour */
    while (stack_pop(&s, &r) == Ok);
    assert(stack_pop(&s, &r) == StackEmpty);
    assert(stack_empty(&s));
    assert(stack_length(&s) == 0);

    if (stack_destroy(&s) != Ok) {
        printf("Error: stack_destroy did not destroy properly...");
        return 1;
    }

    {
        char  buf[32];
        Stack s2;
        if (stack_init_raw(&s2, (uint8_t *)buf, 32, sizeof(int)) != Ok) {
            printf("Error: stack_init_raw did not init properly...");
            return 1;
        }

        assert(stack_capacity(&s2) == 8);
        assert(stack_length(&s2) == 0);
        assert(stack_empty(&s2));

        /* Make sure to not call stack_destroy or free here, since buf[] is static */

        uint8_t buf2[64];
        stack_resize_raw(&s2, buf2, 64);

        assert(stack_capacity(&s2) == 16);
        assert(stack_length(&s2) == 0);
        assert(stack_empty(&s2));
    }

    {
        uint8_t buf[32];
        Stack   st;
        if (stack_init_raw(&st, (uint8_t *)buf, 32, sizeof(int)) != Ok) {
            printf("Error: stack_init_raw did not init properly...");
            return 1;
        }

        for (int i = 1; stack_push(&st, &i) == Ok; i++);

        assert(stack_full(&st));
        assert(stack_capacity(&st) == 8);
        assert(stack_length(&st) == 8);
        assert(!stack_empty(&st));

        /* Make sure to not call stack_destroy or free here, since buf[] is static */

        uint8_t buf2[64];
        stack_resize_raw(&st, buf2, 64);

        assert(stack_capacity(&st) == 16);
        assert(stack_length(&st) == 8);
        assert(!stack_empty(&st));
        assert(!stack_full(&st));

        int r;
        for (int i = 8; i >= 0; i--) {
            if (stack_pop(&st, &r) != Ok)
                return 1;

            printf("%d == %d\n", i, r);
            assert(r == i);
        }
    }
}