#include <assert.h>
#include <inttypes.h>
#include <math.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

#include "expr.h"
#include "finalize.h"
#include "gen.h"
#include "lexer.h"
#include "parser.h"
#include "sym.h"
#include "tokenkind.h"

// set output for code generation

static FILE *out;

void
setParserOut(FILE *out_)
{
    out = out_;
    genSetOutput(out);
}

// for log support

static FILE *logOut;

void
setParserLog(FILE *logOut_)
{
    logOut = logOut_;
}

// error handling
static void
expectedError(const char *expectedStr)
{
    fprintf(stderr, "%zu.%zu: error expected '%s' got '%s'\n", token.pos.line,
            token.pos.col, expectedStr, strTokenKind(token.kind));
    finalizeExit(1);
}

static void
errorAtPos(struct TokenPos pos, const char *msg)
{
    fprintf(stderr, "%zu.%zu: %s\n", pos.line, pos.col, msg);
    finalizeExit(1);
}

static void
expected(enum TokenKind tokenKind)
{
    if (tokenKind != token.kind) {
        expectedError(strTokenKind(tokenKind));
    }
}

// for parsing statements
static bool parseStatement(void);
static bool parseIOHackStatement(void);
static bool parseExprStatement(void);
static bool parseCompoundStatement(void);
static bool parseWhileStatement(void);
static bool parseDoWhileStatement(void);
static bool parseForStatement(void);
static bool parseIfStatement(void);

// for parsing expressions
static const struct Expr *parseAssignmentExpr(void);
static const struct Expr *parseLeftAssocBinaryExpr(int prec);
static const struct Expr *parseUnaryExpr(void);
static const struct Expr *parsePrimaryExpr(void);

void
parse(void)
{
    while (token.kind != EOI) {
        if (!parseStatement()) {
            expectedError("statement");
        }
        deleteAllExpr();
    }
}

//-- for parsing statements

static bool
parseStatement(void)
{
    if (parseIOHackStatement()) {
        return true;
    } else if (parseExprStatement()) {
        return true;
    } else if (parseCompoundStatement()) {
        return true;
    } else if (parseWhileStatement()) {
        return true;
    } else if (parseDoWhileStatement()) {
        return true;
    } else if (parseForStatement()) {
        return true;
    } else if (parseIfStatement()) {
        return true;
    } else {
        return false;
    }
}

static bool
parseIOHackStatement(void)
{
    if (token.kind != DOLLAR) {
        return false;
    }
    getToken();
    if (token.kind == GREATER) {
        getToken();
        // read unsigned integer
        GenReg dest = genGetReg(), val = genGetReg();
        struct TokenPos pos = token.pos;
        const struct Expr *expr = parseAssignmentExpr();
        if (!expr) {
            expectedError("non-empty expression");
        }
        if (!isLValueExpr(expr)) {
            errorAtPos(pos, "L-value expected");
        }
        genInHack(val);
        loadExprAddr(expr, dest);
        genStore(val, dest);
        genUngetReg(dest);
        genUngetReg(val);
    } else if (token.kind == LESS) {
        getToken();
        // print unsigned integer
        GenReg src = genGetReg();
        const struct Expr *expr = parseAssignmentExpr();
        if (!expr) {
            expectedError("non-empty expression");
        }
        loadExpr(expr, src);
        genOutHack(src);
        genUngetReg(src);
    } else {
        expectedError("'>' or '<'");
    }
    expected(SEMICOLON);
    getToken();
    return true;
}

static bool
parseExprStatement(void)
{
    const struct Expr *expr = parseAssignmentExpr();
    if (expr) {
        const struct Expr *folded = constFoldExpr(expr);
        GenReg dest = genGetReg();
        if (logOut) {
            fprintf(logOut, "\\section{}\n");
            printExprTree(expr, logOut);
            if (folded != expr) {
                fprintf(logOut, "after constant folding:\n");
                printExprTree(folded, logOut);
            }
            genSetOutput(logOut);
            fprintf(logOut, "\\begin{lstlisting}\n");
            loadExpr(folded, dest);
            fprintf(logOut, "\\end{lstlisting}\n");
            genSetOutput(out);
        }
        loadExpr(folded, dest);
        genUngetReg(dest);
    } else if (token.kind != SEMICOLON) {
        return false;
    }
    expected(SEMICOLON);
    getToken();
    return true;
}

static bool
parseCompoundStatement(void)
{
    if (token.kind != LBRACE) {
        return false;
    }
    getToken();
    while (parseStatement()) {
    }
    expected(RBRACE);
    getToken();
    return true;
}

static bool
parseWhileStatement(void)
{
    if (token.kind != WHILE) {
        return false;
    }
    getToken();
    expected(LPAREN);
    getToken();

    const struct Expr *cond = parseAssignmentExpr();
    if (!cond) {
        expectedError("assignment expression");
    }

    const char *endLabel = genGetLabel();
    const char *condLabel = genGetLabel();

    genLabelDef(condLabel);

    GenReg dest = genGetReg();
    condJmpExpr(cond, dest, 0, endLabel);
    genUngetReg(dest);

    expected(RPAREN);
    getToken();
    if (!parseCompoundStatement()) {
        expectedError("compound statement");
    }

    genJmp(condLabel);
    genLabelDef(endLabel);
    return true;
}

static bool
parseDoWhileStatement(void)
{
    if (token.kind != DO) {
        return false;
    }
    getToken();

    const char *beginLabel = genGetLabel();

    genLabelDef(beginLabel);
    if (!parseCompoundStatement()) {
        expectedError("compound statement");
    }

    expected(WHILE);
    getToken();

    expected(LPAREN);
    getToken();
    const struct Expr *cond = parseAssignmentExpr();
    if (!cond) {
        expectedError("assignment expression");
    }
    GenReg dest = genGetReg();
    condJmpExpr(cond, dest, beginLabel, 0);
    genUngetReg(dest);

    expected(RPAREN);
    getToken();
    expected(SEMICOLON);
    getToken();
    return true;
}

static bool
parseForStatement(void)
{
    if (token.kind != FOR) {
        return false;
    }
    getToken();

    expected(LPAREN);
    getToken();
    const struct Expr *init = parseAssignmentExpr();
    expected(SEMICOLON);
    getToken();
    const struct Expr *cond = parseAssignmentExpr();
    expected(SEMICOLON);
    getToken();
    const struct Expr *update = parseAssignmentExpr();
    expected(RPAREN);
    getToken();

    if (init) {
        GenReg dest = genGetReg();
        loadExpr(init, dest);
        genUngetReg(dest);
    }

    const char *condLabel = genGetLabel();
    const char *endLabel = genGetLabel();

    genLabelDef(condLabel);
    if (cond) {
        GenReg dest = genGetReg();
        condJmpExpr(cond, dest, 0, endLabel);
        genUngetReg(dest);
    }

    parseCompoundStatement();

    if (update) {
        GenReg dest = genGetReg();
        loadExpr(update, dest);
        genUngetReg(dest);
    }

    genJmp(condLabel);
    genLabelDef(endLabel);
    return true;
}

static bool
parseIfStatement(void)
{
    if (token.kind != IF) {
        return false;
    }
    getToken();

    expected(LPAREN);
    getToken();

    const char *elseLabel = genGetLabel();
    const char *endLabel = genGetLabel();

    const struct Expr *cond = parseAssignmentExpr();
    if (!cond) {
        expectedError("assignment expression");
    }
    GenReg dest = genGetReg();
    condJmpExpr(cond, dest, 0, elseLabel);
    genUngetReg(dest);

    expected(RPAREN);
    getToken();

    if (!parseCompoundStatement()) {
        expectedError("compound statement");
    }

    if (token.kind == ELSE) {
        genJmp(endLabel);
        genLabelDef(elseLabel);
        getToken();
        if (!parseCompoundStatement() && !parseIfStatement()) {
            expectedError("compound statement or if statement");
        }
    } else {
        genLabelDef(elseLabel);
    }
    genLabelDef(endLabel);
    return true;
}

//-- for parsing expressions

static const struct Expr *
parseAssignmentExpr(void)
{
    struct TokenPos pos = token.pos;
    const struct Expr *expr = parseLeftAssocBinaryExpr(1);
    while (token.kind == EQUAL) {
        if (!isLValueExpr(expr)) {
            // instead of many error functions we need a one error handling
            // function that is more flexible to use -> CBE about ellipse
            errorAtPos(pos, "L-value expected");
        }
        getToken();
        const struct Expr *exprRight = parseAssignmentExpr();
        if (!exprRight) {
            expectedError("non-empty expression");
        }
        expr = newBinaryExpr(EK_ASSIGN, expr, exprRight);
    }
    return expr;
}

/*
   Returns 0 if kind is not a left associative binary operator.
   Otherwise returns a precedence > 0
*/
static int
tokenKindPrec(enum TokenKind kind)
{
    switch (kind) {
        case ASTERISK:
        case SLASH:
        case PERCENT:
            return 13;
        case PLUS:
        case MINUS:
            return 12;
        case LESS:
        case LESS_EQUAL:
        case GREATER:
        case GREATER_EQUAL:
            return 11;
        case EQUAL2:
        case NOT_EQUAL:
            return 10;
        default:
            return 0;
    }
}

enum ExprKind
makeBinaryExprKind(enum TokenKind kind)
{
    switch (kind) {
        case ASTERISK:
            return EK_MUL;
        case SLASH:
            return EK_DIV;
        case PERCENT:
            return EK_MOD;
        case PLUS:
            return EK_ADD;
        case MINUS:
            return EK_SUB;
        case LESS:
            return EK_LESS;
        case LESS_EQUAL:
            return EK_LESS_EQUAL;
        case GREATER:
            return EK_GREATER;
        case GREATER_EQUAL:
            return EK_GREATER_EQUAL;
        case EQUAL2:
            return EK_EQUAL;
        case NOT_EQUAL:
            return EK_NOT_EQUAL;
        default:
            printf("kind: %s (%d)\n", strTokenKind(kind), kind);
            assert(0);
            return 0;
    }
}

static const struct Expr *
parseLeftAssocBinaryExpr(int prec)
{
    const struct Expr *expr = parseUnaryExpr();
    for (int p = tokenKindPrec(token.kind); p >= prec; --p) {
        while (tokenKindPrec(token.kind) == p) {
            enum ExprKind op = makeBinaryExprKind(token.kind);
            getToken();
            const struct Expr *rightExpr = parseLeftAssocBinaryExpr(p + 1);
            if (!rightExpr) {
                expectedError("non-empty expression");
            }
            expr = newBinaryExpr(op, expr, rightExpr);
        }
    }
    return expr;
}

static const struct Expr *
parseUnaryExpr(void)
{
    if (token.kind == PLUS || token.kind == MINUS) {
        enum TokenKind op = token.kind;
        getToken();
        const struct Expr *expr = parseUnaryExpr();
        if (!expr) {
            expectedError("non-empty expression");
        }
        if (op == MINUS) {
            return newUnaryExpr(EK_UNARY_MINUS, expr);
        }
        return newUnaryExpr(EK_UNARY_PLUS, expr);
    }
    return parsePrimaryExpr();
}

static const struct Expr *
parsePrimaryExpr(void)
{
    if (token.kind == IDENTIFIER) {
        const struct UStr *identifier = UStrAdd(token.val.cstr);
        struct Sym *sym = SymFind(identifier);
        if (!sym) {
            SymAdd(identifier);
        }
        getToken();
        return newIdentifierExpr(identifier);
    } else if (token.kind == DEC_LITERAL) {
        uint64_t uint = strtoull(token.val.cstr, 0, 10);
        getToken();
        return newUnsignedLiteralExpr(uint);
    } else if (token.kind == HEX_LITERAL) {
        uint64_t uint = strtoull(token.val.cstr, 0, 10);
        getToken();
        return newUnsignedLiteralExpr(uint);
    } else if (token.kind == OCT_LITERAL) {
        uint64_t uint = strtoull(token.val.cstr, 0, 10);
        getToken();
        return newUnsignedLiteralExpr(uint);
    } else if (token.kind == LPAREN) {
        getToken();
        const struct Expr *expr = parseAssignmentExpr();
        expected(RPAREN);
        getToken();
        return expr;
    }
    return 0;
}