#pragma once

#include <optional>
#include <vector>
#include <cassert>

#include "parser.hxx"

class Generator {
public:
    inline explicit Generator(node::Prog prog) : m_prog(std::move(prog)) {
        // Erm!!!
    }

    void gen_scope(const node::Scope* scope) {
        begin_scope();
        for (const node::Stmt* stmt : scope->stmts) {
            gen_stmt(stmt);
        }
        end_scope();
    }

    void gen_stmt(const node::Stmt* stmt) {
        struct StmtVisitor { // Visiting the homies
            Generator& gen;
            void operator()(const node::StmtExit* stmt_exit) const {
                gen.gen_expr(stmt_exit->expr);
                gen.m_output << "    mov rax, 60\n";
                gen.pop("rdi");
                gen.m_output << "    syscall\n";
            }

            void operator()(const node::StmtVar* stmt_var) {
                if (std::ranges::find_if(gen.m_vars.cbegin(), gen.m_vars.cend(), [&](const Var& var){return var.name == stmt_var->ident.value.value();}) != gen.m_vars.cend()) {
                    err::redefine(stmt_var->ident.value.value());
                    exit(EXIT_FAILURE);
                }
                gen.m_vars.push_back({.name = stmt_var->ident.value.value(), .stack_loc = gen.m_stack_size});
                gen.gen_expr(stmt_var->expr);
            }

            void operator()(const node::StmtIf* stmt_if) const {
                gen.gen_expr(stmt_if->expr);
                gen.pop("rax");
                const std::string label = gen.create_label();
                gen.m_output << "    test rax, rax\n";
                gen.m_output << "    jz " << label << "\n";
                gen.gen_scope(stmt_if->scope);
                gen.m_output << "\n" << label << ":\n";
            }

            void operator()(const node::Scope* scope) {
                gen.gen_scope(scope);
            }
        };

        StmtVisitor visitor {.gen = *this};
        std::visit(visitor, stmt->var);
    }

    void gen_term(const node::Term* term) {
        struct TermVisitor { 
            Generator& gen;

            void operator()(const node::TermIntLit* term_int_lit) const {
                gen.m_output << "    mov rax, " << term_int_lit->int_lit.value.value() << "\n";
                gen.push("rax");
            }

            void operator()(const node::TermIdent* term_ident) const {
                const auto it = std::ranges::find_if(gen.m_vars.cbegin(), gen.m_vars.cend(), [&](const Var& var){return var.name == term_ident->ident.value.value();});
                if (it == gen.m_vars.cend()) {
                    err::undeclared(term_ident->ident.value.value());
                    exit(EXIT_FAILURE);
                }
                std::stringstream offset;
                offset << "QWORD [rsp + " << (gen.m_stack_size - it->stack_loc - 1) * 8 << "]";
                gen.push(offset.str());
            }
            void operator()(const node::TermParen* term_paren) const {
                gen.gen_expr(term_paren->expr);
            }
        };
        TermVisitor visitor ({.gen = *this});
        std::visit(visitor, term->var);
    }

    void gen_bin_expr(const node::BinExpr* bin_expr) {
        struct BinExprVisitor {
            Generator& gen;
            void operator()(const node::BinExprAdd* add) {
                gen.gen_expr(add->rhs);
                gen.gen_expr(add->lhs);
                gen.pop("rax");
                gen.pop("rbx");
                gen.m_output << "    add rax, rbx\n";
                gen.push("rax");
            }

            void operator()(const node::BinExprMul* mul) {
                gen.gen_expr(mul->rhs);
                gen.gen_expr(mul->lhs);
                gen.pop("rax");
                gen.pop("rbx");
                gen.m_output << "    mul rbx\n";
                gen.push("rax");
            }

            void operator()(const node::BinExprSub* sub) {
                gen.gen_expr(sub->rhs);
                gen.gen_expr(sub->lhs);
                gen.pop("rax");
                gen.pop("rbx");
                gen.m_output << "    sub rax, rbx\n";
                gen.push("rax");
            }

            void operator()(const node::BinExprDiv* div) {
                gen.gen_expr(div->rhs);
                gen.gen_expr(div->lhs);
                gen.pop("rax");
                gen.pop("rbx");
                gen.m_output << "    div rbx\n";
                gen.push("rax");
            }
        };

        BinExprVisitor visitor {.gen = *this};
        std::visit(visitor, bin_expr->var);
    }

    void gen_expr(const node::Expr* expr) {
        struct ExprVisitor {
            Generator& gen;

            void operator()(const node::Term* term) const {
                gen.gen_term(term);
            }

            void operator()(const node::BinExpr* bin_expr) const {
                gen.gen_bin_expr(bin_expr);
            }
        };

        ExprVisitor visitor {.gen = *this};
        std::visit(visitor, expr->var);
    }

    [[nodiscard]] inline std::string gen_prog() {
        m_output << "global _start\n_start:\n";

        for (const node::Stmt* stmt : m_prog.stmts) {
            gen_stmt(stmt);
        }

        m_output << "    mov rax, 60\n";
        m_output << "    mov rdi, 0\n";
        m_output << "    syscall\n";
        return m_output.str();
    }
private:

    void begin_scope() {
        m_scopes.push_back(m_vars.size());
    }

    void end_scope(){
        const size_t pop_count = m_vars.size() - m_scopes.back();
        m_output << "    add rsp, " << pop_count * 8 << "\n";
        m_stack_size -= pop_count;
        for (int i = 0; i < pop_count; i++) {
            m_vars.pop_back();
        }
        m_scopes.pop_back();
    }

    std::string create_label() {
        return ".label" + std::to_string(m_label_count++);
    }

    void push(const std::string& reg) {
        m_output << "    push " << reg << "\n";
        m_stack_size++;
    }

    void pop(const std::string& reg) {
        m_output << "    pop " << reg << "\n";
        m_stack_size--;
    }

    struct Var {
        std::string name;
        size_t stack_loc;
    };

    const node::Prog m_prog;
    std::stringstream m_output;
    size_t m_stack_size = 0;
    std::vector<Var> m_vars {};
    std::vector<size_t> m_scopes {};
    int m_label_count = 0;
};