#pragma once

#include <optional>
#include <vector>
#include <variant>

#include "tokenization.hxx"
#include "error.hxx"
#include "arena.hxx"

namespace node {
    struct Expr;

    struct TermIntLit {
        Token int_lit;
    };

    struct TermIdent {
        Token ident;
    };

    struct TermParen {
        Expr* expr;
    };

    struct BinExprAdd {
        Expr* lhs;
        Expr* rhs;
    };

    struct BinExprMul {
        Expr* lhs;
        Expr* rhs;
    };

    struct BinExprSub {
        Expr* lhs;
        Expr* rhs;
    };

    struct BinExprDiv {
        Expr* lhs;
        Expr* rhs;
    };

    struct BinExpr {
        std::variant<BinExprAdd*, BinExprMul*, BinExprSub*, BinExprDiv*> var;
    };

    struct Term {
        std::variant<TermIntLit*, TermIdent*, TermParen*> var;
    };

    struct Expr {
        std::variant<Term*, BinExpr*> var;
    };

    struct StmtExit {
        Expr* expr;
    };

    struct StmtVar {
        Token ident;
        Expr* expr;
    };

    struct Stmt {
        std::variant<StmtExit*, StmtVar*> var;
    };

    struct Prog {
        std::vector<Stmt*> stmts;
    };
}

class Parser {
public:
    inline explicit Parser(std::vector<Token> tokens) 
    : m_tokens(std::move(tokens)),
    m_allocator(1024 * 1024 * 4) { // 4 MB should be ok..?
        // Erm!
    }

    std::optional<node::Term*> parse_term() {
        if (auto int_lit = try_consume(TokenType::int_lit)) {
            auto term_int_lit = m_allocator.alloc<node::TermIntLit>();
            term_int_lit->int_lit = int_lit.value();
            auto term = m_allocator.alloc<node::Term>();
            term->var = term_int_lit;
            return term;
        } else if(auto ident = try_consume(TokenType::ident)) {
            auto term_ident = m_allocator.alloc<node::TermIdent>();
            term_ident->ident = ident.value();
            auto term = m_allocator.alloc<node::Term>();
            term->var = term_ident;
            return term;
        } else if(auto open_paren = try_consume(TokenType::open_paren)) {
            auto expr = parse_expr();
            if (!expr.has_value()) {
                err::expression({});
                exit(EXIT_FAILURE);
            }
            try_consume(TokenType::close_paren, ")");
            auto term_paren = m_allocator.alloc<node::TermParen>();
            term_paren->expr = expr.value();
            auto term = m_allocator.alloc<node::Term>();
            term->var = term_paren;
            return term;
        } else {
            return {};
        }
    }

    std::optional<node::Expr*> parse_expr(int min_prec = 0) {
        std::optional<node::Term*> term_lhs = parse_term();
        if (!term_lhs.has_value()) {
            return {};
        }

        auto expr_lhs = m_allocator.alloc<node::Expr>();
        expr_lhs->var = term_lhs.value();

        while (true) {
            std::optional<Token> curr_tok = peek();
            std::optional<int> prec;
            if (curr_tok.has_value()) {
                prec = bin_prec(curr_tok->type);
                if (!prec.has_value() || prec < min_prec) {
                    break; // breaking bad
                }
            } else {
                break; // breaking bad 2
            }

            Token op = consume();
            int next_min_prec = prec.value() + 1;
            auto expr_rhs = parse_expr(next_min_prec);
            if (!expr_rhs.has_value()) {
                err::expression({});
                exit(EXIT_FAILURE);
            }

            auto expr = m_allocator.alloc<node::BinExpr>();
            auto expr_lhs2 = m_allocator.alloc<node::Expr>(); // LHS two: The Sequel.
            if (op.type == TokenType::plus) {
                auto add = m_allocator.alloc<node::BinExprAdd>();
                expr_lhs->var = term_lhs.value();
                expr_lhs2->var = expr_lhs->var;
                add->lhs = expr_lhs2;
                add->rhs = expr_rhs.value();
                expr->var = add;
            } else if (op.type == TokenType::star) {
                auto mul = m_allocator.alloc<node::BinExprMul>();
                expr_lhs->var = term_lhs.value();
                expr_lhs2->var = expr_lhs->var;
                mul->lhs = expr_lhs2;
                mul->rhs = expr_rhs.value();
                expr->var = mul;
            } else if (op.type == TokenType::minus) {
                auto sub = m_allocator.alloc<node::BinExprSub>();
                expr_lhs->var = term_lhs.value();
                expr_lhs2->var = expr_lhs->var;
                sub->lhs = expr_lhs2;
                sub->rhs = expr_rhs.value();
                expr->var = sub;
            } else if (op.type == TokenType::slash) {
                auto div = m_allocator.alloc<node::BinExprDiv>();
                expr_lhs->var = term_lhs.value();
                expr_lhs2->var = expr_lhs->var;
                div->lhs = expr_lhs2;
                div->rhs = expr_rhs.value();
                expr->var = div;
            } else {
                err::syntax({});
                exit(EXIT_FAILURE);
            }
            expr_lhs->var = expr;
        }
        return expr_lhs;
    }

    std::optional<node::Stmt*> parse_stmt() {
        if (peek().has_value() && peek().value().type == TokenType::exit && peek(1).has_value() && peek(1).value().type == TokenType::open_paren) {
            consume();
            consume();
            auto stmt_exit = m_allocator.alloc<node::StmtExit>();
            if (auto node_expr = parse_expr()) {
                stmt_exit->expr = node_expr.value();
            } else {
                err::expression({});
                exit(EXIT_FAILURE);
            }

            try_consume(TokenType::close_paren, ")");
            try_consume(TokenType::semi, ";");

            auto stmt = m_allocator.alloc<node::Stmt>();
            stmt->var = stmt_exit;
            return stmt;
        } else if (peek().has_value() && peek().value().type == TokenType::var 
        && peek(1).has_value() && peek(1).value().type == TokenType::ident 
        && peek(2).has_value() && peek(2).value().type == TokenType::equals) { // Holy else-if statement
            consume();
            auto stmt_var = m_allocator.alloc<node::StmtVar>();
            stmt_var->ident = consume();
            consume();
            if (auto expr = parse_expr()) {
                stmt_var->expr = expr.value();
            } else {
                err::expression({});
            }
            try_consume(TokenType::semi, ";");

            auto stmt = m_allocator.alloc<node::Stmt>();
            stmt->var = stmt_var;
            return stmt;
        } else {
            return {};
        }
    } 

    std::optional<node::Prog> parse_prog() {
        node::Prog prog;
        while (peek().has_value()) {
            if (auto stmt = parse_stmt()) {
                prog.stmts.push_back(stmt.value());
            } else {
                err::statement({});
                exit(EXIT_FAILURE);
            }
        }
        return prog;
    }

private:
    [[nodiscard]] inline std::optional<Token> peek(int offset = 0) const {
        if (m_index + offset >= m_tokens.size()) {
            return {};
        } else {
            return m_tokens.at(m_index + offset);
        }
    }

    inline Token consume() {
        return m_tokens.at(m_index++);
    }

    inline Token try_consume(TokenType type, const std::string& err) {
        if (peek().has_value() && peek().value().type == type) {
            return consume();
        } else {
            err::syntax(err);
            exit(EXIT_FAILURE);
        }
    }

    inline std::optional<Token> try_consume(TokenType type) {
        if (peek().has_value() && peek().value().type == type) {
            return consume();
        } else {
            return {};
        }
    }

    const std::vector<Token> m_tokens;
    size_t m_index = 0;
    ArenaAllocator m_allocator;
};