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

// Structures defined in this file
typedef struct BlockHeader BlockHeader;
typedef struct Closure Closure;
typedef struct Frame Frame;
typedef struct Text Text;
typedef struct UnionType UnionType;
typedef struct Var Var;

// You can change this number to decide how much memory the program is allowed
//   0 - 193 bytes   -> The program refuses to start
// 194 - 217 bytes   -> The program starts but doesn't do anything
// 218 +     bytes   -> The program runs successfully
#define MEM_SIZE 193   //1024 * 1024 // 1 MB

#define EXIT_GRACEFULLY 0
#define EXIT_MEMORY_LIMIT_TOO_LOW 11
#define EXIT_COULD_NOT_INIT_MEMORY 12
#define EXIT_OUT_OF_MEMORY_DURING_INITIALIZATION 13

#define EXIT_IMPOSSIBLE_FUNCTION_PHASE 77
#define EXIT_IMPOSSIBLE_FUNCTION_ID 78
#define EXIT_IMPOSSIBLE_FUNCTION_NO_CONTINUATION 79

#define EXIT_IF_NULL(ptr) ({                                                   \
    if (!ptr) {                                                                \
        free(memory);                                                          \
        printf("Failed to initialize program: out of memory\n");               \
        exit(EXIT_OUT_OF_MEMORY_DURING_INITIALIZATION);                        \
    }                                                                          \
})

/* ============================================================================
 *                                  MEMORY
 * ============================================================================
 */

// Memory block structure
// This block functions as a header for memory allocations on the heap
typedef struct BlockHeader {
    int free;
    struct BlockHeader *next;
} BlockHeader;

#define BLOCK_SIZE(header) ({                                                  \
    size_t s;                                                                  \
    if (header->next) {                                                        \
        s = (uint8_t *)(header->next) - (uint8_t *)header;                     \
    } else {                                                                   \
        s = (uint8_t *)(stack_top) - (uint8_t *)header;                        \
    }                                                                          \
    s - sizeof(BlockHeader);                                                   \
})

// Global memory pool
static uint8_t *memory = NULL;
static uintptr_t heap_base;
static uintptr_t heap_top;
static uintptr_t stack_base;
static uintptr_t stack_top;

// Initialize the custom memory manager
void init_memory() {
    if (MEM_SIZE < sizeof(BlockHeader) + 1) {
        printf("Memory limit is too low!\n");
        exit(EXIT_MEMORY_LIMIT_TOO_LOW);
    }

    memory = (uint8_t *)malloc(MEM_SIZE);
    if (!memory) {
        printf("Memory allocation failed!\n");
        exit(EXIT_COULD_NOT_INIT_MEMORY);
    }
    heap_base = (uintptr_t)memory;
    heap_top = heap_base;
    stack_base = (uintptr_t)(memory + MEM_SIZE);
    stack_top = stack_base;

    BlockHeader *curr = (BlockHeader *)heap_base;
    curr->free = 1;
    curr->next = NULL;
}

// Macro for allocating a number of bytes on the heap
#define HEAP_ALLOC(nbytes) ({                                                  \
    void *ptr = NULL;                                                          \
    BlockHeader *curr = (BlockHeader *)heap_base;                              \
    while(curr) {                                                              \
        size_t size = (BLOCK_SIZE(curr));                                      \
        if (curr->free && size >= (nbytes)) {                                  \
            /* If there's enough memory left over, split off into two */       \
            if (size >= (nbytes) + sizeof(BlockHeader) + 1) {                  \
                BlockHeader *new_block = (BlockHeader *)(                      \
                    (uint8_t *)curr + sizeof(BlockHeader) + (nbytes)           \
                );                                                             \
                new_block->free = 1;                                           \
                new_block->next = curr->next;                                  \
                curr->next = new_block;                                        \
                if ((uint8_t *)heap_top < (uint8_t *)new_block) {                                    \
                    heap_top = (uintptr_t)new_block + sizeof(BlockHeader);                           \
                }                                                              \
            }                                                                  \
            /* Assign block to new memory */                                   \
            curr->free = 0;                                                    \
            ptr = (uint8_t *)curr + sizeof(BlockHeader);                       \
            break;                                                             \
        }                                                                      \
        /* Otherwise, go to the next block */                                  \
        curr = curr->next;                                                     \
    }                                                                          \
    ptr;                                                                       \
})

// Macro for freeing bytes at a given address on the heap
#define HEAP_FREE(ptr) ({                                                      \
    if (ptr) {                                                                 \
        BlockHeader *block = (BlockHeader *)(                                  \
            (uint8_t *)(ptr) - sizeof(BlockHeader)                             \
        );                                                                     \
        block->free = 1;                                                       \
        /* Merge adjacent free blocks */                                       \
        BlockHeader *curr = (BlockHeader *)heap_base;                          \
        while (curr) {                                                         \
            if (curr->free && curr->next && curr->next->free) {                \
                curr->size += sizeof(BlockHeader) + curr->next->size;          \
                curr->next = curr->next->next;                                 \
            } else {                                                           \
                curr = curr->next;                                             \
            }                                                                  \
        }                                                                      \
    }                                                                          \
})

// Pointer for debugging: print the heap to stdout
#define HEAP_PRINT_WIDTH 60
#define HEAP_PRINT() ({                                                        \
    BlockHeader *curr = (BlockHeader *)heap_base;                              \
    while (curr) {                                                             \
        size_t size = BLOCK_SIZE(curr);                                        \
        printf("BlockHeader at address %ld (physical %p): %lu byte",           \
            (uint8_t *)curr - memory, curr, size                               \
        );                                                                     \
        if (size == 1) { printf(" of "); } else { printf("s of "); }           \
        if (curr->free) {                                                      \
            printf("free memory\n    ");                                       \
        } else {                                                               \
            printf("reserved memory\n    ");                                   \
        }                                                                      \
        unsigned char *content = (uint8_t *)curr;                              \
        for (int i = 0; i < size + sizeof(BlockHeader); i++) {                 \
            if (i == sizeof(BlockHeader)) { printf("| "); }                    \
            printf("%02X ", content[i]);                                       \
            if (i > HEAP_PRINT_WIDTH) { printf("..."); break; }                \
        }                                                                      \
        printf("\n    ");                                                      \
        for (int i = 0; i < size + sizeof(BlockHeader); i++) {                 \
            if (i == sizeof(BlockHeader)) { printf("| "); }                    \
            if (content[i] == 0) { printf("   ");                              \
            } else if (content[i] >= 32 && content[i] < 127) {                 \
                printf("%c  ", content[i]);                                    \
            } else { printf("__ "); }                                          \
            if (i > HEAP_PRINT_WIDTH) { printf("..."); break; }                              \
        }                                                                      \
        printf("\n");                                                          \
        curr = curr->next;                                                     \
    }                                                                          \
})

#define SPACE_FILL(n) for (int i=0; i<n; i++) { printf(" "); }

// Macro for allocating a number of bytes on the stack
#define STACK_ALLOC(nbytes) ({                                                 \
    void *ptr = NULL;                                                          \
    if (stack_top >= heap_top + (nbytes)) {                                    \
        stack_top -= (nbytes);                                                 \
        ptr = (void *)stack_top;                                               \
    }                                                                          \
    ptr;                                                                       \
})

#define STACK_PRINT_VARIABLE_MAX 5
#define STACK_PRINT(topPtr) ({                                                 \
    Frame *cursor = topPtr;                                                    \
    while (cursor) {                                                           \
        Frame *next = cursor->cont;                                            \
        size_t size;                                                           \
        if (next) {                                                            \
            size = (uint8_t *)(next) - (uint8_t *)(cursor);                    \
        } else {                                                               \
            size = (uint8_t *)(stack_base) - (uint8_t *)(cursor);              \
        }                                                                      \
        printf("Stack layer at offset -%ld (phsyical %p): %ld byte",           \
            (uint8_t *)(stack_base) - (uint8_t *)(cursor), cursor, size        \
        );                                                                     \
        if (size == 1) { printf("\n    "); } else { printf("s\n    "); }       \
        unsigned char *content = (uint8_t *)cursor;                            \
        printf("| Code ID"); SPACE_FILL(sizeof(CodeID)*3 - 7 - 1);             \
        printf("| Phase"); SPACE_FILL(sizeof(int)*3 - 5 - 2);                  \
        printf("| Prev out"); SPACE_FILL(sizeof(Var *)*3-8 - 2);               \
        printf("| Continuation"); SPACE_FILL(sizeof(Frame *)*3-12 - 1);        \
        for (int i = 0; i < size - sizeof(Frame); i++) {                       \
            printf("| Variable %d", i + 2); SPACE_FILL(2);                     \
            if (i + 2 == STACK_PRINT_VARIABLE_MAX) { break; }                  \
        }                                                                      \
        printf("|\n    | ");                                                   \
        for (int i = 0; i < sizeof(Frame); i++) {                              \
            printf("%02X ", content[i]);                                       \
        }                                                                      \
        for (int i = sizeof(Frame); i < size; i++) {                           \
            if ((i - sizeof(Frame)) % sizeof(Var) == 0) { printf("|"); }       \
            printf("%02X ", content[i]);                                       \
            if (i - sizeof(Frame) == STACK_PRINT_VARIABLE_MAX * sizeof(Var)) { \
                break;                                                         \
            }                                                                  \
        }                                                                      \
        printf("|\n");                                                         \
        cursor = cursor->cont;                                                 \
    }                                                                          \
})
// #define STACK_PRINT ({                                                         \
//     printf("Stack from base %p to top %p:\n", (void *)stack_base, (void *)stack_top); \
//     uint8_t *curr = (uint8_t *)stack_top;                                      \
//     while ((uintptr_t)curr < stack_base) {                                     \
//         printf("Address %ld (physical %p): \n    ", curr - memory, curr);             \
//         for (int i = 0; i < STACK_PRINT_WIDTH && (uintptr_t)(curr + i) < stack_base; i++) { \
//             printf("%02X ", curr[i]);                                          \
//         }                                                                      \
//         printf("\n    ");                                                      \
//         for (int i = 0; i < STACK_PRINT_WIDTH && (uintptr_t)(curr + i) < stack_base; i++) { \
//             if (curr[i] == 0) {                                                \
//                 printf("   ");                                                 \
//             } else if (curr[i] >= 32 && curr[i] < 127) {                       \
//                 printf("%c  ", curr[i]);                                       \
//             } else {                                                           \
//                 printf("__ ");                                                 \
//             }                                                                  \
//         }                                                                      \
//         printf("\n");                                                          \
//         curr += STACK_PRINT_WIDTH;                                             \
//     }                                                                          \
// })


// Macro for freeing a number of bytes on the stack
#define STACK_FREE(nbytes) ({                                                  \
    stack_top += (nbytes);                                                     \
})

#define MEM_DIAGNOSTICS(curr) ({                                               \
    printf("----------------------------\n");                                  \
    printf("Heap base at address %ld (physical %p)\n",                         \
        (uint8_t *)heap_base - (uint8_t *)memory, (uint8_t *)heap_base                    \
    );                                                                         \
    printf("Heap top at address %ld (physical %p)\n",                          \
        (uint8_t *)heap_top - (uint8_t *)memory, (uint8_t *)heap_top                      \
    );                                                                         \
    HEAP_PRINT();                                                              \
    printf("----------------------------\n");                                  \
    printf("Heap base at address %ld (physical %p)\n",                         \
        (uint8_t *)heap_base - (uint8_t *)memory, (uint8_t *)heap_base                    \
    );                                                                         \
    printf("Heap top at address %ld (physical %p)\n",                          \
        (uint8_t *)heap_top - (uint8_t *)memory, (uint8_t *)heap_top                      \
    );                                                                         \
    printf("Stack top at address %ld (physical %p)\n",                         \
        (uint8_t *)stack_top - (uint8_t *)memory, (uint8_t *)stack_top                    \
    );                                                                         \
    printf("Stack base at address %ld (physical %p)\n",                        \
        (uint8_t *)stack_base - (uint8_t *)memory, (uint8_t *)stack_base                  \
    );                                                                         \
    printf("----------------------------\n");                                  \
    if (curr) {                                                                \
        STACK_PRINT(curr);                                                     \
    } else {                                                                   \
        printf("Stack is NULL!\n");                                            \
    }                                                                          \
    printf("Stack top at address %ld (physical %p)\n",                         \
        (uint8_t *)stack_top - (uint8_t *)memory, (uint8_t *)stack_top                    \
    );                                                                         \
    printf("Stack base at address %ld (physical %p)\n",                        \
        (uint8_t *)stack_base - (uint8_t *)memory, (uint8_t *)stack_base                  \
    );                                                                         \
    printf("----------------------------\n");                                  \
})


/* ============================================================================
 *                            VARIABLES & VALUES
 * ============================================================================
 */

// Types of values that can be stored in the heap
typedef union {
    int i;
    char c;
    Closure *f;
    Text *t;
    UnionType *u;
} Value;

#define HEAP_ALLOC_VALUE() HEAP_ALLOC(sizeof(Value))

// Types of variables stored in a frame
typedef enum {
    ORIGINAL_VARIABLE,
    MUTABLE_REFERENCE,
    IMMUTABLE_REFERENCE
} VarType;

// Variable
typedef struct Var {
    VarType var_type;
    Value *value;
} Var;

#define HEAP_ALLOC_VARIABLE() HEAP_ALLOC(sizeof(Var))

// Names of functions
typedef enum {
    CODE_F,
    CODE_G,
    CODE_H,
    CODE_EXIT  // Indicates termination
} CodeID;
 
// A closure can represent a function that has yet to gather all inputs
typedef struct Closure {
    CodeID code_id;
    Var variables[];
} Closure;

// A Text struct represents a string
typedef struct Text {
    size_t len;
    char s[];
} Text;

// A Union type represents a type that can be one of several values.
// For example:
//      type Maybe a = Just a | Nothing
typedef struct UnionType {
    int name;
    Value values[];
} UnionType;

#define HEAP_ALLOC_CLOSURE(code_id, arity) ({                                  \
    Closure *c = HEAP_ALLOC(sizeof(Closure) + arity * sizeof(Var));            \
    if (c) {                                                                   \
        c->code_id = code_id;                                                  \
    }                                                                          \
    c;                                                                         \
})

#define HEAP_ALLOC_TEXT(length) ({                                                \
    Text *t = HEAP_ALLOC(sizeof(Text) + length * sizeof(char));                \
    if (t) {                                                                   \
        t->len = length;                                                       \
        for (int i = 0; i < length; i++) {                                     \
            t->s[i] = '\0';                                                    \
        }                                                                      \
    }                                                                          \
    t;                                                                         \
})

#define HEAP_ALLOC_UNION_TYPE(name, arity) ({                                  \
    UnionType *u = HEAP_ALLOC(sizeof(UnionType) + arity * sizeof(Value));      \
    if (u) {                                                                   \
        u->name = name;                                                        \
    }                                                                          \
    u;                                                                         \
})

// Function arity
#define ARITY_F 1
#define ARITY_G 1
#define ARITY_H 1
#define ARITY_EXIT 1

typedef struct Frame {
    // Function to execute
    CodeID code_id;
    
    // Phase of the function
    int phase;
    
    // Return variable from previous frame
    Var *prevOut;
    
    // CPS-style next frame to execute
    struct Frame *cont;
    
    // Previously captured variables
    Var variables[];
} Frame;

#define FRAME_SIZE(arity) sizeof(Frame) + (arity - 1) * sizeof(Var)

#define STACK_ALLOC_FRAME(code_name, arity) ({                                 \
    Frame *f = STACK_ALLOC(FRAME_SIZE(arity));                                 \
    if(f) {                                                                    \
        f->code_id = code_name;                                                \
        f->phase = 0;                                                          \
        f->prevOut = NULL;                                                     \
        f->cont = NULL;                                                        \
    }                                                                          \
    f;                                                                         \
})

#define STACK_FREE_FRAME(arity) STACK_FREE(FRAME_SIZE(arity))

int main(void) {
    init_memory();
    
    Text *t = HEAP_ALLOC_TEXT(50);                          EXIT_IF_NULL(t);
    Value *v = HEAP_ALLOC_VALUE();                          EXIT_IF_NULL(v);
    Var *var = HEAP_ALLOC_VARIABLE();                       EXIT_IF_NULL(var);
    Frame *fnsh = STACK_ALLOC_FRAME(CODE_EXIT, ARITY_EXIT); EXIT_IF_NULL(fnsh);
    Frame *curr = STACK_ALLOC_FRAME(CODE_H, ARITY_H);       EXIT_IF_NULL(curr);
    
    sprintf(t->s, "Kaokkokos shall prevail by the hands of Alkbaard!");
    
    v->t = t;
    var->var_type = ORIGINAL_VARIABLE;
    var->value = v;
    curr->prevOut = var;
    curr->cont = fnsh;
    
    // STACK_PRINT(curr);
    
    while (curr) {
        // MEM_DIAGNOSTICS(curr);
        switch (curr->code_id) {
            case CODE_F: {
                // f(x) = print(x) and return x
                Var *x = curr->prevOut;
                Frame *next = curr->cont;
                STACK_FREE_FRAME(ARITY_F);
    
                if (x) {
                    printf("%s\n", x->value->t->s);
                } else {
                    printf("NULL: Out of memory!\n");
                }
    
                curr = next;
                break;
            }
            case CODE_G: {
                // g(x) = String.upper(x)
                Var *x = curr->prevOut;
                Frame *next = curr->cont;
                STACK_FREE_FRAME(ARITY_G);

                if (x) {
                    for (int i = 0; i < x->value->t->len; i++) {
                        char c = x->value->t->s[i];
                        if (c >= 'a' && c <= 'z') {
                            x->value->t->s[i] = c - ('a' - 'A');
                        }
                    }

                    next->prevOut = x;
                }

                curr = next;
                break;
            }
            case CODE_H : {
                // h(x) = f(g(x))
                Var *x = curr->prevOut;
                Frame *next = curr->cont;
                STACK_FREE_FRAME(ARITY_H);

                if (x) {
                    Frame *f = STACK_ALLOC_FRAME(CODE_F, ARITY_F);
                    
                    if (f) {
                        Frame *g = STACK_ALLOC_FRAME(CODE_G, ARITY_G);

                        if (g) {
                            f->cont = next;

                            g->prevOut = x;
                            g->cont = f;
                            curr = g;
                        } else {
                            // No memory available! Free f & exit this function
                            STACK_FREE_FRAME(ARITY_F);
                            next->prevOut = NULL;
                            curr = next;
                        }
                    } else {
                        // No memory available! Exit this function
                        next->prevOut = NULL;
                        curr = next;
                    }
                } else {
                    // prevOut is NULL - we were out of memory!
                    curr = next;
                }

                break;
            }
            case CODE_EXIT: {
                // HEAP_PRINT();
                free(memory);
                exit(EXIT_GRACEFULLY);
                break;
            }
            default: {
                free(memory);
                printf("ERROR: Reached impossible function id!\n");
                exit(EXIT_IMPOSSIBLE_FUNCTION_ID);
                break;
            }
        }
    }

    free(memory);
    
    printf("ERROR: Reached function with no continuation!\n");

    return EXIT_IMPOSSIBLE_FUNCTION_NO_CONTINUATION;
}