#include "Compiler.h"

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

#include "Code.h"
#include "Memory.h"

typedef struct Variable Variable;

struct Variable {
    Variable* next;
    i32 index;
    char name[];
};

typedef struct Function Function;

struct Function {
    Function* next;
    i32 address;
    i32 valueAddress;
    i32 arguments;
    char name[];
};

typedef struct Return Return;

struct Return {
    Return* next;
    i32 address;
    bool hasArgument;
};

typedef struct {
    Tokenizer* tokenizer;
    Code* code;
    Error error;
    int line;
    jmp_buf jump;
    Variable* variables;
    Function* functions;
    Return* returns;
    bool inFunction;
} Context;

#define THROW_ERROR(format, ...)                                   \
    snprintf(                                                      \
        c->error.text, sizeof(c->error.text), "Line %d | " format, \
        c->line __VA_OPT__(, ) __VA_ARGS__);                       \
    longjmp(c->jump, 1)

[[noreturn]] static void throwOutOfCodeMemory(Context* c) {
    THROW_ERROR("Too less memory for code");
}

static void codePushInstruction(Context* c, Instruction i) {
    if(bufferWriteU8(&c->code->code, i)) {
        throwOutOfCodeMemory(c);
    }
}

static void codePushI32(Context* c, i32 i) {
    if(bufferWriteI32(&c->code->code, i)) {
        throwOutOfCodeMemory(c);
    }
}

static void codePushConstantString(Context* c, const char* s) {
    if(bufferWriteString(&c->code->code, s)) {
        throwOutOfCodeMemory(c);
    }
}

static size_t codeGetWritePosition(const Context* c) {
    return c->code->code.writeIndex;
}

static size_t codePushInstructionI32(Context* c, Instruction i, i32 v) {
    codePushInstruction(c, i);
    size_t pos = codeGetWritePosition(c);
    codePushI32(c, v);
    return pos;
}

static void codeRewriteI32(Context* c, size_t pos, i32 i) {
    size_t oldPos = codeGetWritePosition(c);
    c->code->code.writeIndex = pos;
    codePushI32(c, i);
    c->code->code.writeIndex = oldPos;
}

static bool hasLocalVar(const Context* c) {
    return c->variables != nullptr && c->variables->index < 0;
}

static i32 addVariable(Context* c, const char* name, bool forceGlobal) {
    for(Variable* v = c->variables; v != nullptr; v = v->next) {
        if(!forceGlobal && c->inFunction && v->index >= 0) {
            break;
        }
        if(strcmp(v->name, name) == 0) {
            return v->index;
        }
    }
    size_t l = strlen(name) + 1;
    Variable* v = memoryAllocate(sizeof(Variable) + l);
    if(v == nullptr) {
        THROW_ERROR("Too less memory for variables");
    }
    v->next = c->variables;
    if(c->inFunction) {
        v->index = hasLocalVar(c) ? c->variables->index - 1 : -1;
    } else {
        v->index = c->variables != nullptr ? c->variables->index + 1 : 0;
    }
    memcpy(v->name, name, l);
    c->variables = v;
    return v->index;
}

static void addRawFunction(
    Context* c, const char* name, i32 address, i32 valueAddress,
    i32 arguments) {
    size_t l = strlen(name) + 1;
    Function* f = memoryAllocate(sizeof(Function) + l);
    if(f == nullptr) {
        THROW_ERROR("Too less memory for functions");
    }
    f->next = c->functions;
    f->address = address;
    f->valueAddress = valueAddress;
    f->arguments = arguments;
    memcpy(f->name, name, l);
    c->functions = f;
}

static void addFunction(
    Context* c, const char* name, i32 address, i32 arguments) {
    for(Function* f = c->functions; f != nullptr; f = f->next) {
        if(strcmp(f->name, name) != 0) {
            continue;
        } else if(f->address >= 0) {
            THROW_ERROR("Function '%s' registered again", name);
        } else if(f->arguments != arguments) {
            THROW_ERROR("Function '%s' with inconsistent arguments", name);
        }
        f->address = address;
        codeRewriteI32(c, (size_t)f->valueAddress, address);
    }
    addRawFunction(c, name, address, -1, arguments);
}

static i32 addCallFunction(
    Context* c, const char* name, i32 address, i32 arguments) {
    for(Function* f = c->functions; f != nullptr; f = f->next) {
        if(strcmp(f->name, name) == 0 && f->address >= 0) {
            if(f->arguments != arguments) {
                THROW_ERROR("Function '%s' with inconsistent arguments", name);
            }
            return f->address;
        }
    }
    addRawFunction(c, name, -1, address, arguments);
    return -1;
}

static void checkForInvalidFunctions(Context* c) {
    for(Function* f = c->functions; f != nullptr; f = f->next) {
        if(f->address < 0) {
            THROW_ERROR("Unmapped function call '%s'", f->name);
        }
    }
}

static void addReturn(Context* c, i32 address, bool hasArgument) {
    if(c->returns != nullptr && c->returns->hasArgument != hasArgument) {
        THROW_ERROR("Inconsistent returns in function");
    }
    Return* r = memoryAllocate(sizeof(Return));
    if(r == nullptr) {
        THROW_ERROR("Too less memory for returns");
    }
    r->next = c->returns;
    r->hasArgument = hasArgument;
    r->address = address;
    c->returns = r;
}

#define CLEAN_LIST(Type, field)   \
    do {                          \
        Type* v = c->field;       \
        while(v != nullptr) {     \
            Type* next = v->next; \
            memoryFree(v);        \
            v = next;             \
        }                         \
        c->field = nullptr;       \
    } while(false)

static void cleanReturns(Context* c) {
    CLEAN_LIST(Return, returns);
}

static void cleanContext(Context* c) {
    CLEAN_LIST(Variable, variables);
    CLEAN_LIST(Function, functions);
    cleanReturns(c);
}

[[noreturn]] static void unexpectedToken(Context* c, Token token, int line) {
    char buffer[128];
    tokenizerPrintToken(&token, buffer, sizeof(buffer));
    THROW_ERROR("Unexpected %s token %d", buffer, line);
}

static Token consumeToken(Context* c, TokenType type) {
    Token actual = tokenizerNext(c->tokenizer);
    if(actual.type != type) {
        unexpectedToken(c, actual, __LINE__);
    }
    return actual;
}

static bool peekToken(Context* c, TokenType type) {
    return tokenizerPeek(c->tokenizer).type == type;
}

static bool consumeTokenIf(Context* c, TokenType type) {
    if(peekToken(c, type)) {
        consumeToken(c, type);
        return true;
    }
    return false;
}

static void consumeNewline(Context* c) {
    if(peekToken(c, TT_INVALID)) {
        return;
    }
    consumeToken(c, TT_NEWLINE);
    c->line++;
}

static void compileReadVariable(Context* c, Token token, bool forceGlobal) {
    i32 index = addVariable(c, token.stringValue, forceGlobal);
    codePushInstructionI32(c, READ_VARIABLE, index);
}

static void compileCallFunction(Context* c, Token t);

static void compileConstant(Context* c) {
    Token token = tokenizerNext(c->tokenizer);
    if(token.type == TT_STRING) {
        codePushInstruction(c, PUSH_CONSTANT_STRING);
        codePushConstantString(c, token.stringValue);
    } else if(token.type == TT_INT32) {
        codePushInstructionI32(c, PUSH_INT32, token.intValue);
    } else if(token.type == TT_LITERAL) {
        if(peekToken(c, TT_OPEN_ROUND_BRACKET)) {
            compileCallFunction(c, token);
        } else {
            compileReadVariable(c, token, false);
        }
    } else if(token.type == TT_DOLLAR) {
        token = consumeToken(c, TT_LITERAL);
        compileReadVariable(c, token, true);
    } else {
        unexpectedToken(c, token, __LINE__);
    }
}

static void compileUnary(Context* c) {
    if(consumeTokenIf(c, TT_NOT)) {
        compileUnary(c);
        codePushInstruction(c, NOT);
    } else {
        compileConstant(c);
    }
}

static void compileMul(Context* c) {
    compileUnary(c);
    while(true) {
        if(consumeTokenIf(c, TT_MULTIPLY)) {
            compileUnary(c);
            codePushInstruction(c, MUL);
        } else if(consumeTokenIf(c, TT_DIVIDE)) {
            compileUnary(c);
            codePushInstruction(c, DIV);
        } else {
            break;
        }
    }
}

static void compileAdd(Context* c) {
    compileMul(c);
    while(true) {
        if(consumeTokenIf(c, TT_ADD)) {
            compileMul(c);
            codePushInstruction(c, ADD);
        } else if(consumeTokenIf(c, TT_SUBTRACT)) {
            compileMul(c);
            codePushInstruction(c, SUB);
        } else {
            break;
        }
    }
}

static void compileCompare(Context* c) {
    compileAdd(c);
    while(true) {
        if(consumeTokenIf(c, TT_EQUAL)) {
            compileAdd(c);
            codePushInstruction(c, EQUAL);
        } else if(consumeTokenIf(c, TT_NOT_EQUAL)) {
            compileAdd(c);
            codePushInstruction(c, NOT_EQUAL);
        } else if(consumeTokenIf(c, TT_GREATER)) {
            compileAdd(c);
            codePushInstruction(c, GREATER);
        } else if(consumeTokenIf(c, TT_SMALLER)) {
            compileAdd(c);
            codePushInstruction(c, SMALLER);
        } else if(consumeTokenIf(c, TT_GREATER_OR_EQUAL)) {
            compileAdd(c);
            codePushInstruction(c, GREATER_OR_EQUAL);
        } else if(consumeTokenIf(c, TT_SMALLER_OR_EQUAL)) {
            compileAdd(c);
            codePushInstruction(c, SMALLER_OR_EQUAL);
        } else {
            break;
        }
    }
}

static void compileLogical(Context* c) {
    compileCompare(c);
    while(true) {
        if(consumeTokenIf(c, TT_AND)) {
            compileCompare(c);
            codePushInstruction(c, AND);
        } else if(consumeTokenIf(c, TT_OR)) {
            compileCompare(c);
            codePushInstruction(c, OR);
        } else {
            break;
        }
    }
}

static void compileExpression(Context* c) {
    compileLogical(c);
}

static void compileLine(Context* c, Token token);

static void compileIf(Context* c) {
    compileExpression(c);
    size_t posIndex = codePushInstructionI32(c, JUMP_ON_0, 0);
    consumeNewline(c);
    while(!peekToken(c, TT_END)) {
        compileLine(c, tokenizerNext(c->tokenizer));
    }
    consumeToken(c, TT_END);
    consumeNewline(c);
    codeRewriteI32(c, posIndex, (i32)codeGetWritePosition(c));
}

static void compileSetVariable(Context* c, const char* name, bool forceGlobal) {
    consumeToken(c, TT_ASSIGN);
    compileExpression(c);
    codePushInstructionI32(c, SET_VARIABLE, addVariable(c, name, forceGlobal));
}

static i32 compileFunctionArguments(Context* c) {
    consumeToken(c, TT_OPEN_ROUND_BRACKET);
    i32 vars = 0;
    while(!peekToken(c, TT_CLOSE_ROUND_BRACKET)) {
        if(vars > 0) {
            consumeToken(c, TT_COMMA);
        }
        Token varToken = consumeToken(c, TT_LITERAL);
        addVariable(c, varToken.stringValue, false);
        vars++;
    }
    consumeToken(c, TT_CLOSE_ROUND_BRACKET);
    return vars;
}

static void rewriteReturns(Context* c, i32 address) {
    for(Return* r = c->returns; r != nullptr; r = r->next) {
        codeRewriteI32(c, (size_t)r->address, address);
    }
    cleanReturns(c);
}

static void compileFunction(Context* c) {
    if(c->inFunction) {
        THROW_ERROR("Functions in functions are not allowed");
    }
    c->inFunction = true;
    size_t functionEnd = codePushInstructionI32(c, JUMP, 0);
    i32 functionStart = (i32)codeGetWritePosition(c);
    Token t = consumeToken(c, TT_LITERAL);
    size_t stackVarsLoc = codePushInstructionI32(c, PUSH_STACK_VARIABLES, 0);
    Variable* resetVar = c->variables;

    i32 vars = compileFunctionArguments(c);
    addFunction(c, t.stringValue, functionStart, vars);
    while(!peekToken(c, TT_END)) {
        compileLine(c, tokenizerNext(c->tokenizer));
    }
    consumeToken(c, TT_END);
    i32 popVars = hasLocalVar(c) ? -c->variables->index : 0;
    codePushInstructionI32(c, PUSH_INT32, 0);
    codePushInstructionI32(c, RETURN_VALUE, popVars);
    rewriteReturns(c, popVars);
    codeRewriteI32(c, stackVarsLoc, popVars - vars);

    while(c->variables != resetVar) {
        Variable* next = c->variables->next;
        memoryFree(c->variables);
        c->variables = next;
    }

    codeRewriteI32(c, functionEnd, (i32)codeGetWritePosition(c));
    c->inFunction = false;
}

static void compileCallFunction(Context* c, Token t) {
    // used to store return address and variable index
    codePushInstructionI32(c, PUSH_INT32, 0);
    codePushInstructionI32(c, PUSH_INT32, 0);

    consumeToken(c, TT_OPEN_ROUND_BRACKET);
    i32 offset = 0;
    while(!peekToken(c, TT_CLOSE_ROUND_BRACKET)) {
        if(offset > 0) {
            consumeToken(c, TT_COMMA);
        }
        offset++;
        compileExpression(c);
    }
    consumeToken(c, TT_CLOSE_ROUND_BRACKET);

    codePushInstruction(c, JUMP_SUB);
    size_t pos = codeGetWritePosition(c);
    codePushI32(c, addCallFunction(c, t.stringValue, (i32)pos, offset));
    codePushI32(c, offset);
}

static void compileReturn(Context* c) {
    if(peekToken(c, TT_NEWLINE)) {
        consumeNewline(c);
        codePushInstructionI32(c, PUSH_INT32, 0);
    } else {
        compileExpression(c);
    }
    i32 address = (i32)codePushInstructionI32(c, RETURN_VALUE, 0);
    addReturn(c, address, true);
}

static void compileLineLiteral(Context* c, Token token) {
    const char* s = token.stringValue;
    if(strcmp(s, "print") == 0) {
        while(!peekToken(c, TT_NEWLINE)) {
            compileExpression(c);
            codePushInstruction(c, PRINT);
        }
        consumeNewline(c);
        codePushInstruction(c, PRINT_NEWLINE);
    } else if(peekToken(c, TT_ASSIGN)) {
        compileSetVariable(c, s, false);
    } else if(peekToken(c, TT_OPEN_ROUND_BRACKET)) {
        compileCallFunction(c, token);
    } else {
        THROW_ERROR("Unexpected literal(%s)", s);
    }
}

static void compileLine(Context* c, Token token) {
    if(token.type == TT_NEWLINE) {
        c->line++;
    } else if(token.type == TT_IF) {
        compileIf(c);
    } else if(token.type == TT_FUNCTION) {
        compileFunction(c);
    } else if(token.type == TT_RETURN) {
        compileReturn(c);
    } else if(token.type == TT_DOLLAR) {
        token = consumeToken(c, TT_LITERAL);
        compileSetVariable(c, token.stringValue, true);
    } else if(token.type != TT_LITERAL) {
        unexpectedToken(c, token, __LINE__);
    } else {
        compileLineLiteral(c, token);
    }
}

static void parseTokens(Context* c) {
    codeReset(c->code);
    size_t stackVarsLoc = codePushInstructionI32(c, PUSH_STACK_VARIABLES, 0);

    while(true) {
        Token token = tokenizerNext(c->tokenizer);
        if(token.type == TT_INVALID) {
            break;
        }
        compileLine(c, token);
    }

    if(c->variables != nullptr) {
        codeRewriteI32(c, stackVarsLoc, c->variables->index + 1);
    }
    checkForInvalidFunctions(c);
}

Error compileFile(Tokenizer* t, Code* code) {
    Context c = {.tokenizer = t, .code = code, .line = 1};
    if(!setjmp(c.jump)) {
        parseTokens(&c);
    }
    cleanContext(&c);
    return c.error;
}