#ifndef AST_H
#define AST_H

#include "codegen.h"

#include <string>
#include <vector>
#include <memory>

namespace AST {
    class Node {
    public:
        virtual std::string formatted() = 0;
        virtual ~Node() = default;
    };

    class Expression : public Node {
    public:
        virtual llvm::Value* codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
    };

    class Type {
    public:
        virtual ~Type() = default;
        virtual std::string formatted() = 0;
        virtual llvm::Type* codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
    };

    class Statement : public Node {
    public:
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
    };

    class IntLiteralExpression : public Expression {
    private:
        int m_value;
    public:
        IntLiteralExpression(int value) : m_value{ value } {}
        virtual ~IntLiteralExpression() override = default;
        virtual std::string formatted() override;
        virtual llvm::Value* codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };

    class ValueReferenceExpression : public Expression {
    private:
        std::string m_name;
    public:
        ValueReferenceExpression(std::string name) : m_name{ name } {}
        virtual ~ValueReferenceExpression() override = default;
        virtual std::string formatted() override;
        virtual llvm::Value* codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };


    class ReturnStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        ReturnStatement(std::unique_ptr<Expression> expr) : m_expr{ std::move(expr) } {}
        virtual ~ReturnStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };

    class InitializationStatement : public Statement {
    private:
        std::unique_ptr<Type> m_type;
        std::string m_name;
        std::optional<std::unique_ptr<Expression>> m_expr;
    public:
        InitializationStatement(std::unique_ptr<Type> ty, std::string name, std::optional<std::unique_ptr<Expression>> expr)
            : m_type{ std::move(ty) }, m_name{ name }, m_expr{ std::move(expr) } {
        }
        virtual ~InitializationStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };

    class TopLevelStatement : public Node {
    public:
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
    };

    class Function : public TopLevelStatement {
    private:
        std::unique_ptr<Type> m_return_ty;
        std::vector<std::pair<std::string, std::unique_ptr<Type>>> m_params;
        std::string m_name;
        std::vector<std::unique_ptr<Statement>> m_statements;
    public:
        Function(
            std::unique_ptr<Type> return_ty,
            std::vector<std::pair<std::string, std::unique_ptr<Type>>> params,
            std::string name,
            std::vector<std::unique_ptr<Statement>> statements)
            : m_return_ty{ std::move(return_ty) }
            , m_params{ std::move(params) }
            , m_name{ name }
            , m_statements{ std::move(statements) } {
        }
        virtual ~Function() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };

    enum class FundamentalTypeKind {
        Int,
    };

    class FundamentalType : public Type {
    private:
        FundamentalTypeKind m_ty;
    public:
        FundamentalType(FundamentalTypeKind kind) : m_ty{ kind } {}
        virtual ~FundamentalType() override = default;
        virtual std::string formatted() override;
        virtual llvm::Type* codegen(codegen::Builder& builder, codegen::Scope& scope) override;
    };
}

#endif