#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 {
        cAddInt(f->address);
        cAddInt(arguments);
        if(f->returns && noReturn) {
            cAddOperation(OP_POP);
        }
    }
}

static void cLiteral() {
    const char* literal = cReadString();
    if(cConsumeTokenIf(T_OPEN_BRACKET)) {
        cCallFunction(literal, false);
    } 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 cMul() {
    cPrimary();
    while(cConsumeTokenIf(T_MUL)) {
        cPrimary();
        cAddOperation(OP_MUL);
    }
}

static void cAdd() {
    cMul();
    while(cConsumeTokenIf(T_ADD)) {
        cMul();
        cAddOperation(OP_ADD);
    }
}

static void cExpression() {
    cAdd();
}

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;
        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 cFunctionInnerBody(int arguments) {
    int p = cAddPush(arguments);
    int oldLine = line;
    while(!cConsumeTokenIf(T_CLOSE_CURVED_BRACKET)) {
        Token t = cReadTokenAndLine();
        if(t == T_END) {
            cError("unexpected end of file: function not closed on line %d", oldLine);
        }
        cLine(t);
    }
    cAddPop(p, vars[1].entries);
    for(int i = 0; i < returnIndex; i++) {
        cSetInt(returns[i], vars[1].entries);
    }
    returnIndex = 0;
}

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

    int address = code->length;
    returnState = 0;
    cFunctionInnerBody(arguments);
    if(!fmAdd(&functions, name, arguments, address, returnState == 2)) {
        cError("function registered twice on line %d", line);
    }

    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 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;
        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);
        }
        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;
}