#pragma once
#include "prelude/string.hpp"

enum class NodeType
{
    Expression = 0,
    // TODO: abstract "Literal" node type
    IntLiteral,
    Extern,
    FnDecl,
    FnCall,
    Variable,
    VarDecl,
    Compound,
    Program,
    COUNT_NODES,
};

#define NODE_TYPE(x) \
    NodeType GetType() const override { return NodeType::x; }

class Node
{
public:
    virtual NodeType GetType() const = 0;
    virtual ~Node() {}
};

class ExpressionNode : public Node
{
public:
    enum class Operator
    {
        Plus = 0,
        Minus,
        Divide,
        Multiply,
        COUNT_OPERATORS,
    };
public:
    ExpressionNode(Node* left, Node* right, Operator op)
        : m_left(left), m_right(right), m_op(op) {}
    ~ExpressionNode() override {
        delete m_left;
        delete m_right;
    }

    NODE_TYPE(Expression)
public:
    Node* left() const { return m_left; }
    Node* right() const { return m_right; }
    Operator op() const { return m_op; }
private:
    Node* m_left;
    Node* m_right;
    Operator m_op;
};

// TODO: Maybe just LiteralNode with double or int literals support
class IntLiteralNode : public Node
{
public:
    IntLiteralNode(long value)
        : m_int_value(value) {}
    ~IntLiteralNode() override {}

    NODE_TYPE(IntLiteral)
public:
    long integer() const { return m_int_value; }
private:
    long m_int_value;
};

class ExternNode : public Node
{
public:
    // TODO: support multiple extern symbols
    ExternNode(StringView symbol)
        : m_symbol(symbol) {}
    ~ExternNode() override = default;

    NODE_TYPE(Extern)
public:
    const StringView& symbol() const { return m_symbol; }
private:
    StringView m_symbol;
};

class CompoundNode : public Node
{
public:
    CompoundNode() = default;

    NODE_TYPE(Compound)

    // --- Iteration ---
    auto begin() { return m_nodes.begin(); }
    auto end()   { return m_nodes.end();   }
    auto begin() const { return m_nodes.begin(); }
    auto end()   const { return m_nodes.end();   }

    // --- Access by index ---
    Node* operator[](size_t i) { return m_nodes[i]; }
    const Node* operator[](size_t i) const { return m_nodes[i]; }

    Node* at(size_t i) { return m_nodes.at(i); }
    const Node* at(size_t i) const { return m_nodes.at(i); }

    // --- Modifiers ---
    void addNode(Node* n) { m_nodes.push_back(n); }
    void removeNode(size_t idx) { m_nodes.erase(m_nodes.begin() + idx); }

    // If you want full expose for iteration but not modification
    const std::vector<Node*>& nodes() const { return m_nodes; }

    size_t size() const { return m_nodes.size(); }
    bool empty() const { return m_nodes.empty(); }
private:
    std::vector<Node*> m_nodes;
};

class FnDeclNode : public Node
{
public:
    // TODO: support parameters
    FnDeclNode(const StringView& name, CompoundNode* body, View<StringView> params)
        : m_name(name), m_body(body), m_params(params) {}
    ~FnDeclNode() override {
        delete m_body;
    }

    NODE_TYPE(FnDecl)
public:
    const StringView& name() const { return m_name; }
    const CompoundNode* body() const { return m_body; }
    const View<StringView>& params() const { return m_params; }
private:
    StringView m_name;
    CompoundNode* m_body;
    View<StringView> m_params;
};

class FnCallNode : public Node
{
public:
    // TODO: support multiple arguments
    FnCallNode(const StringView& name, View<Node*>&& arg)
        : m_name(name), m_args(arg) {}
    ~FnCallNode() override = default;

    NODE_TYPE(FnCall)
public:
    const StringView& name() const { return m_name; }
    const View<Node*>& args() const { return m_args; }
private:
    StringView m_name;
    View<Node*> m_args;
};

class VariableNode : public Node
{
public:
    VariableNode(const StringView& name)
        : m_name(name) {}
    ~VariableNode() override = default;

    NODE_TYPE(Variable)
public:
    const StringView& name() const { return m_name; }
private:
    StringView m_name;
};

class VarDeclNode : public Node
{
public:
    VarDeclNode(const StringView& name, Node* value)
        : m_name(name), m_value(value) {}
    ~VarDeclNode() override {
        delete m_value;
    }

    NODE_TYPE(VarDecl)
public:
    const StringView& name() const { return m_name; }
    const Node* value() const { return m_value; }
private:
    StringView m_name;
    Node* m_value;
};

class ProgramNode : public Node
{
public:
    ProgramNode() = default;

    NODE_TYPE(Program)
public:
    void PushFunction(FnDeclNode* fn)
    {
        m_funcs.push_back(fn);
    }

    void PushExtern(ExternNode* extrn)
    {
        m_externs.push_back(extrn);
    }
public:
    const std::vector<ExternNode*> externs() const { return m_externs; }
    const std::vector<FnDeclNode*> funcs() const { return m_funcs; }
private:
    std::vector<FnDeclNode*> m_funcs;
    std::vector<ExternNode*> m_externs;
};