#include <setjmp.h>
#include <stdarg.h>
#include <stdio.h>

#include "Compiler.h"
#include "FunctionMap.h"
#include "Operation.h"
#include "StringIntMap.h"
#include "Tokenizer.h"

#define ERROR_LENGTH 256
#define RETURN_BUFFER 16

static jmp_buf errorJump;
static char error[ERROR_LENGTH] = {'\0'};

static ByteCode* code;

static int16 line = 1;

static int varIndex = 0;
static StringIntMap vars[2];
static FunctionMap functions;

static int returns[RETURN_BUFFER];
static int returnIndex = 0;
static int returnState = 0;

static void cError(const char* format, ...) {
    va_list args;
    va_start(args, format);
    vsnprintf(error, ERROR_LENGTH, format, args);
    va_end(args);
    longjmp(errorJump, 0);
}

static int cAddVar(const char* var) {
    int index = vars[varIndex].entries;
    simAdd(vars + varIndex, var, &index);
    return index;
}

static void cUnexpectedToken(Token t) {
    cError("unexpected token on line %d: %s", line, tGetTokenName(t));
}

static void cAddOperation(Operation token) {
    unsigned char c = token;
    bcAddBytes(code, &c, 1);
}

static int cReserveInt() {
    return bcReserveBytes(code, sizeof(int));
}

static void cSetInt(int p, int i) {
    bcSetBytes(code, p, &i, sizeof(int));
}

static void cAddInt(int i) {
    bcAddBytes(code, &i, sizeof(int));
}

static void cAddInt16(int16 i) {
    bcAddBytes(code, &i, sizeof(int16));
}

static void cAddFloat(float f) {
    bcAddBytes(code, &f, sizeof(float));
}

static int cAddPush(int offset) {
    cAddOperation(OP_PUSH_VARS);
    int p = cReserveInt();
    cAddInt(offset);
    return p;
}

static void cAddPop(int p, int vars) {
    cAddOperation(OP_POP_VARS);
    cAddInt(vars);
    cSetInt(p, vars);
}

static Token cReadTokenAndLine() {
    Token t = tReadToken();
    if(tReadInt16(&line)) {
        return t;
    }
    return T_END;
}

static void cConsumeToken(Token wanted) {
    Token t = cReadTokenAndLine();
    if(wanted != t) {
        cError("unexpected token on line %d: expected '%s' got '%s'", line, tGetTokenName(wanted), tGetTokenName(t));
    }
}

static bool cConsumeTokenIf(Token t) {
    if(tPeekToken() == t) {
        cReadTokenAndLine();
        return true;
    }
    return false;
}

static void cConstantInt() {
    int value;
    if(!tReadInt(&value)) {
        cError("int token without an int on line %d", line);
    }
    cAddOperation(OP_PUSH_INT);
    cAddInt(value);
}

static void cConstantFloat() {
    float value;
    if(!tReadFloat(&value)) {
        cError("float token without a float on line %d", line);
    }
    cAddOperation(OP_PUSH_FLOAT);
    cAddFloat(value);
}

static const char* cReadString() {
    const char* literal = tReadString();
    if(literal == NULL) {
        cError("literal without string on line %d", line);
    }
    return literal;
}

static void cGetVar(const char* var) {
    cAddOperation(OP_GET);
    cAddInt(cAddVar(var));
}

static void cExpression();

static int cCallFunctionArguments() {
    int arguments = 0;
    while(!cConsumeTokenIf(T_CLOSE_BRACKET)) {
        arguments++;
        cExpression();
        if(cConsumeTokenIf(T_COMMA) && tPeekToken() == T_CLOSE_BRACKET) {
            cUnexpectedToken(tPeekToken());
        }
    }
    return arguments;
}

static void cCallFunction(const char* literal, bool noReturn) {
    cAddOperation(OP_PUSH_INT);
    cAddInt(0);
    int arguments = cCallFunctionArguments();
    Function* f = fmSearch(&functions, literal, arguments);
    cAddOperation(OP_GOSUB);
    if(f == NULL) {
        fmEnqueue(&functions, literal, arguments, line, cReserveInt(), noReturn);
        cAddInt(arguments);
        cAddOperation(OP_NOTHING);
    } else {
        if(!noReturn && !f->returns) {
            cError("function '%s' needs a return value on line %d", f->name, line);
        }
        cAddInt(f->address);
        cAddInt(arguments);
        if(f->returns && noReturn) {
            cAddOperation(OP_POP);
        }
    }
}

static void cPostIncrement(const char* literal) {
    cAddOperation(OP_POST_INCREMENT);
    cAddInt(cAddVar(literal));
}

static void cPostDecrement(const char* literal) {
    cAddOperation(OP_POST_DECREMENT);
    cAddInt(cAddVar(literal));
}

static void cLiteral() {
    const char* literal = cReadString();
    if(cConsumeTokenIf(T_OPEN_BRACKET)) {
        cCallFunction(literal, false);
    } else if(cConsumeTokenIf(T_INCREMENT)) {
        cPostIncrement(literal);
    } else if(cConsumeTokenIf(T_DECREMENT)) {
        cPostDecrement(literal);
    } else {
        cGetVar(literal);
    }
}

static void cPrimary() {
    Token t = cReadTokenAndLine();
    switch(t) {
        case T_INT: cConstantInt(); break;
        case T_FLOAT: cConstantFloat(); break;
        case T_NULL: cAddOperation(OP_PUSH_NULL); break;
        case T_TRUE: cAddOperation(OP_PUSH_TRUE); break;
        case T_FALSE: cAddOperation(OP_PUSH_FALSE); break;
        case T_OPEN_BRACKET:
            cExpression();
            cConsumeToken(T_CLOSE_BRACKET);
            break;
        case T_LITERAL: cLiteral(); break;
        default: cUnexpectedToken(t); break;
    }
}

static void cPreIncrement() {
    cConsumeToken(T_LITERAL);
    cAddOperation(OP_PRE_INCREMENT);
    cAddInt(cAddVar(cReadString()));
}

static void cPreDecrement() {
    cConsumeToken(T_LITERAL);
    cAddOperation(OP_PRE_DECREMENT);
    cAddInt(cAddVar(cReadString()));
}

static void cPreUnary() {
    if(cConsumeTokenIf(T_SUB)) {
        cPrimary();
        cAddOperation(OP_INVERT_SIGN);
    } else if(cConsumeTokenIf(T_INCREMENT)) {
        cPreIncrement();
    } else if(cConsumeTokenIf(T_DECREMENT)) {
        cPreDecrement();
    } else {
        cPrimary();
    }
}

static void cMul() {
    cPreUnary();
    while(true) {
        if(cConsumeTokenIf(T_MUL)) {
            cPreUnary();
            cAddOperation(OP_MUL);
        } else if(cConsumeTokenIf(T_DIV)) {
            cPreUnary();
            cAddOperation(OP_DIV);
        } else if(cConsumeTokenIf(T_MOD)) {
            cPreUnary();
            cAddOperation(OP_MOD);
        } else {
            break;
        }
    }
}

static void cAdd() {
    cMul();
    while(true) {
        if(cConsumeTokenIf(T_ADD)) {
            cMul();
            cAddOperation(OP_ADD);
        } else if(cConsumeTokenIf(T_SUB)) {
            cMul();
            cAddOperation(OP_SUB);
        } else {
            break;
        }
    }
}

static void cComparison() {
    cAdd();
    while(true) {
        if(cConsumeTokenIf(T_LESS)) {
            cAdd();
            cAddOperation(OP_LESS);
        } else if(cConsumeTokenIf(T_LESS_EQUAL)) {
            cAdd();
            cAddOperation(OP_GREATER);
            cAddOperation(OP_NOT);
        } else if(cConsumeTokenIf(T_GREATER)) {
            cAdd();
            cAddOperation(OP_GREATER);
        } else if(cConsumeTokenIf(T_GREATER_EQUAL)) {
            cAdd();
            cAddOperation(OP_LESS);
            cAddOperation(OP_NOT);
        } else if(cConsumeTokenIf(T_EQUAL)) {
            cAdd();
            cAddOperation(OP_EQUAL);
        } else if(cConsumeTokenIf(T_NOT_EQUAL)) {
            cAdd();
            cAddOperation(OP_EQUAL);
            cAddOperation(OP_NOT);
        } else {
            break;
        }
    }
}

static void cLogical() {
    cComparison();
    while(true) {
        if(cConsumeTokenIf(T_AND)) {
            cAddOperation(OP_DUPLICATE);
            cAddOperation(OP_IF_GOTO);
            int p = cReserveInt();
            cComparison();
            cAddOperation(OP_AND);
            cSetInt(p, code->length);
        } else if(cConsumeTokenIf(T_OR)) {
            cAddOperation(OP_DUPLICATE);
            cAddOperation(OP_NOT);
            cAddOperation(OP_IF_GOTO);
            int p = cReserveInt();
            cComparison();
            cAddOperation(OP_OR);
            cSetInt(p, code->length);
        } else {
            break;
        }
    }
}

static void cExpression() {
    cLogical();
}

static void cSetVar(const char* literal) {
    cExpression();
    cConsumeToken(T_SEMICOLON);
    cAddOperation(OP_SET);
    cAddInt(cAddVar(literal));
}

static void cLineLiteral() {
    const char* literal = cReadString();
    Token t = cReadTokenAndLine();
    switch(t) {
        case T_SET: cSetVar(literal); break;
        case T_OPEN_BRACKET:
            cCallFunction(literal, true);
            cConsumeToken(T_SEMICOLON);
            break;
        case T_INCREMENT:
            cPostIncrement(literal);
            cConsumeToken(T_SEMICOLON);
            cAddOperation(OP_POP);
            break;
        case T_DECREMENT:
            cPostDecrement(literal);
            cConsumeToken(T_SEMICOLON);
            cAddOperation(OP_POP);
            break;
        default: cUnexpectedToken(t);
    }
}

static int cFunctionArguments() {
    int arguments = 0;
    while(!cConsumeTokenIf(T_CLOSE_BRACKET)) {
        cConsumeToken(T_LITERAL);
        arguments++;
        cAddVar(cReadString());
        if(cConsumeTokenIf(T_COMMA) && tPeekToken() != T_LITERAL) {
            cUnexpectedToken(tPeekToken());
        }
    }
    return arguments;
}

static void cLine(Token t);

static void cConsumeBody() {
    cConsumeToken(T_OPEN_CURVED_BRACKET);
    int oldLine = line;
    while(!cConsumeTokenIf(T_CLOSE_CURVED_BRACKET)) {
        Token t = cReadTokenAndLine();
        if(t == T_END) {
            cError("unexpected end of file: non closed curved bracket on line %d", oldLine);
        }
        cLine(t);
    }
}

static void cLinkReturns() {
    for(int i = 0; i < returnIndex; i++) {
        cSetInt(returns[i], vars[1].entries);
    }
    returnIndex = 0;
}

static void cFunctionBody(const char* name, int arguments) {
    int oldLine = line;
    cAddOperation(OP_GOTO);
    int gotoIndex = cReserveInt();

    int address = code->length;
    returnState = 0;

    int p = cAddPush(arguments);
    cConsumeBody();
    cAddPop(p, vars[1].entries);

    cLinkReturns();

    if(!fmAdd(&functions, name, arguments, address, returnState == 2)) {
        cError("function registered twice on line %d", oldLine);
    }

    cAddOperation(OP_RETURN);
    cSetInt(gotoIndex, code->length);
}

static void cFunction() {
    if(varIndex == 1) {
        cError("function inside function on line %d", line);
    }
    cConsumeToken(T_LITERAL);
    const char* name = cReadString();
    cConsumeToken(T_OPEN_BRACKET);
    varIndex = 1;
    vars[1].entries = 0;
    cFunctionBody(name, cFunctionArguments());
    varIndex = 0;
}

static void cAddReturn() {
    cAddOperation(OP_POP_VARS);
    returns[returnIndex++] = cReserveInt(vars);
    cAddOperation(OP_RETURN);
}

static void cReturn() {
    if(varIndex == 0) {
        cError("return without a function on line %d", line);
    } else if(returnIndex >= RETURN_BUFFER) {
        cError("too much returns in function around line %d", line);
    }
    if(cConsumeTokenIf(T_SEMICOLON)) {
        if(returnState == 2) {
            cError("mixed return type on line %d", line);
        }
        returnState = 1;
        cAddReturn();
    } else {
        if(returnState == 1) {
            cError("mixed return type on line %d", line);
        }
        returnState = 2;
        cExpression();
        cAddOperation(OP_SET_RETURN);
        cAddReturn();
        cConsumeToken(T_SEMICOLON);
    }
}

static void cPrint() {
    cExpression();
    cConsumeToken(T_SEMICOLON);
    cAddOperation(OP_PRINT);
}

static void cIf() {
    cConsumeToken(T_OPEN_BRACKET);
    cExpression();
    cConsumeToken(T_CLOSE_BRACKET);
    cAddOperation(OP_IF_GOTO);
    int ifP = cReserveInt();
    cConsumeBody();
    cSetInt(ifP, code->length);

    if(cConsumeTokenIf(T_ELSE)) {
        cAddOperation(OP_GOTO);
        int elseP = cReserveInt();
        cSetInt(ifP, code->length);
        if(cConsumeTokenIf(T_IF)) {
            cIf();
        } else {
            cConsumeBody();
        }
        cSetInt(elseP, code->length);
    }
}

static void cWhile() {
    int start = code->length;
    cConsumeToken(T_OPEN_BRACKET);
    cExpression();
    cConsumeToken(T_CLOSE_BRACKET);
    cAddOperation(OP_IF_GOTO);
    int ifP = cReserveInt();
    cConsumeBody();
    cAddOperation(OP_GOTO);
    cAddInt(start);
    cSetInt(ifP, code->length);
}

static void cLine(Token t) {
    cAddOperation(OP_LINE);
    cAddInt16(line);
    switch(t) {
        case T_PRINT: cPrint(); break;
        case T_LITERAL: cLineLiteral(); break;
        case T_FUNCTION: cFunction(); break;
        case T_RETURN: cReturn(); break;
        case T_IF: cIf(); break;
        case T_WHILE: cWhile(); break;
        case T_INCREMENT:
            cPreIncrement();
            cConsumeToken(T_SEMICOLON);
            cAddOperation(OP_POP);
            break;
        case T_DECREMENT:
            cPreDecrement();
            cConsumeToken(T_SEMICOLON);
            cAddOperation(OP_POP);
            break;
        default: cUnexpectedToken(t);
    }
}

static void cForEachLine() {
    Token t = cReadTokenAndLine();
    while(t != T_END) {
        cLine(t);
        t = cReadTokenAndLine();
    }
}

static void cLinkQueuedFunctions() {
    for(int i = 0; i < functions.queueEntries; i++) {
        Function* f = fmSearch(&functions, functions.queue[i].name, functions.queue[i].arguments);
        if(f == NULL) {
            cError("unknown function on line %d", functions.queue[i].line);
        } else if(!functions.queue[i].noReturn && !f->returns) {
            cError("function '%s' needs a return value on line %d", f->name, functions.queue[i].line);
        }
        cSetInt(functions.queue[i].reserved, f->address);
        if(functions.queue[i].noReturn && f->returns) {
            code->code[functions.queue[i].reserved + sizeof(int) * 2] = OP_POP;
        }
    }
}

static void cAllocAndCompile() {
    varIndex = 0;
    returnIndex = 0;
    simInit(vars);
    simInit(vars + 1);
    fmInit(&functions);
    if(!setjmp(errorJump)) {
        int p = cAddPush(0);
        cForEachLine();
        cAddPop(p, vars[varIndex].entries);
        cLinkQueuedFunctions();
    }
    fmDelete(&functions);
    simDelete(vars + 1);
    simDelete(vars);
}

ByteCode* cCompile() {
    error[0] = '\0';
    code = bcInit();
    cAllocAndCompile();
    if(error[0] != '\0') {
        bcDelete(code);
        return NULL;
    }
    return code;
}

const char* cGetError() {
    return error;
}