#ifndef WRAITH_STACK_H
#define WRAITH_STACK_H

#include <stdint.h>
#include "ast.h"

typedef struct StackFrame
{
    Node **statements;
    uint32_t statementCount;
    uint32_t statementCapacity;

    Node **declarations;
    uint32_t declarationCount;
    uint32_t declarationCapacity;
} StackFrame;

typedef struct Stack
{
    StackFrame *stackFrames;
    uint32_t stackCapacity;
    uint32_t stackIndex;
} Stack;

Stack* CreateStack()
{
    Stack *stack = (Stack*) malloc(sizeof(Stack));
    stack->stackCapacity = 4;
    stack->stackFrames = (StackFrame*) malloc(sizeof(StackFrame) * stack->stackCapacity);
    stack->stackIndex = 0;
    return stack;
}

void PushStackFrame(Stack *stack)
{
    stack->stackIndex += 1;

    if (stack->stackIndex == stack->stackCapacity)
    {
        stack->stackCapacity += 1;
        stack->stackFrames = (StackFrame*) realloc(stack->stackFrames, sizeof(StackFrame) * stack->stackCapacity);

        stack->stackFrames[stack->stackIndex].statementCapacity = 0;
        stack->stackFrames[stack->stackIndex].declarationCapacity = 0;
    }

    stack->stackFrames[stack->stackIndex].statementCount = 0;
    stack->stackFrames[stack->stackIndex].declarationCount = 0;
}

void PopStackFrame(Stack *stack)
{
    stack->stackIndex -= 1;
}

void AddStatement(Stack *stack, Node *statementNode)
{
    StackFrame *stackFrame = &stack->stackFrames[stack->stackIndex];
    if (stackFrame->statementCount == stackFrame->statementCapacity)
    {
        stackFrame->statementCapacity += 1;
        stackFrame->statements = (Node**) realloc(stackFrame->statements, stackFrame->statementCapacity);
    }
    stackFrame->statements[stackFrame->statementCount] = statementNode;
    stackFrame->statementCount += 1;
}

Node** GetStatements(Stack *stack, uint32_t *pCount)
{
    *pCount = stack->stackFrames[stack->stackIndex].statementCount;
    return stack->stackFrames[stack->stackIndex].statements;
}

void AddDeclaration(Stack *stack, Node *declarationNode)
{
    StackFrame *stackFrame = &stack->stackFrames[stack->stackIndex];
    if (stackFrame->declarationCount == stackFrame->declarationCapacity)
    {
        stackFrame->declarationCapacity += 1;
        stackFrame->declarations = (Node**) realloc(stackFrame->declarations, stackFrame->declarationCapacity);
    }
    stackFrame->declarations[stackFrame->declarationCount] = declarationNode;
    stackFrame->declarationCount += 1;
}

Node** GetDeclarations(Stack *stack, uint32_t *pCount)
{
    *pCount = stack->stackFrames[stack->stackIndex].declarationCount;
    return stack->stackFrames[stack->stackIndex].declarations;
}

#endif /* WRAITH_STACK_H */