/* File 0 */

namespace LX::AST
{
	// Base node that everything else inherits from
	struct Node
	{
		// Enum for storing the type of node //
		// Used so a pointer to Node can be used and then turned into it's true type //
		enum NodeType
		{
			// General Nodes //

			IDENTIFIER,
			NUMBER_LITERAL,
			OPERATION,

			// Control flow Nodes //

			RETURN_STATEMENT,

			// If an error happened somewhere //
			UNDEFINED = -1
		};

		// Constructor to set the node type //
		Node(NodeType type)
			: m_Type(type)
		{}

		// Virtual destructor because of polymorphism //
		virtual ~Node() = default;

		// Function for generating LLVN IR (Intermediate representation) //
		virtual llvm::Value* GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder) = 0;

		// Function for generating C/C++ code (Currently not implemented) //
		//virtual void GenC() = 0;

		// The type of the node //
		const NodeType m_Type;
	};

	class NumberLiteral : public Node
	{
		public:
			// Constructor to set values and automatically set type
			NumberLiteral(std::string num);

			// Function for generating LLVN IR (Intermediate representation) //
			llvm::Value* GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder) override;

		private:
			// The number it stores
			// Yes the number is stored as a string
			// It's horrible I know
			std::string m_Number;
	};

	//
	class Operation : public Node
	{
		public:
			// Constructor to set values and automatically set type
			Operation(std::unique_ptr<Node> lhs, Token::TokenType op, std::unique_ptr<Node> rhs);

			// Function for generating LLVN IR (Intermediate representation) //
			llvm::Value* GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder) override;

		private:
			// The sides of the operation
			// Unary operations are handled by a different class
			std::unique_ptr<Node> m_Lhs, m_Rhs;

			// The operation to be applied to the two sides
			Token::TokenType m_Operand;
	};

	//
	class ReturnStatement : public Node
	{
		public:
			// Constructor to set values and automatically set type
			ReturnStatement(std::unique_ptr<Node> val);

			// Function for generating LLVN IR (Intermediate representation) //
			llvm::Value* GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder) override;

		private:
			// What it is returning (can be null)
			std::unique_ptr<Node> m_Val;
	};
}

namespace LX
{
	struct IRGenerationError {};

	struct FunctionDefinition
	{
		FunctionDefinition()
			: body{}
		{}
		
		std::vector<std::unique_ptr<AST::Node>> body;
	};

	struct FileAST
	{
		FileAST()
			: functions{}
		{}

		std::vector<FunctionDefinition> functions;
	};

	FileAST TurnTokensIntoAbstractSyntaxTree(std::vector<Token>& tokens, std::ofstream* log);

	void GenerateIR(FileAST& ast);
}

/* File 1 */

#include <Parser.h>

#include <LLVM.h>
#include <Util.h>

namespace LX::AST
{
	llvm::Value* NumberLiteral::GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder)
	{
		// Converts the string to it's int equivalent
		// Will eventually need to do floating point stuff here as well
		int number = std::stoi(m_Number);

		// Returns it as a llvm value (if valid)
		llvm::Value* out = llvm::ConstantInt::get(llvm::Type::getInt32Ty(context), number, true);
		ThrowIf<IRGenerationError>(out == nullptr);
		return out;
	}

	llvm::Value* Operation::GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder)
	{
		// Gets the IR for both sides of the operation
		llvm::Value* lhs = m_Lhs->GenIR(context, module, builder);
		llvm::Value* rhs = m_Rhs->GenIR(context, module, builder);

		// If either side is null then return null to prevent invalid IR //
		if (lhs == nullptr || rhs == nullptr)
		{
			ThrowIf<IRGenerationError>(true);
			return nullptr;
		}

		// Will eventually get the correct operator but for now everything is add
		llvm::Value* out = builder.CreateAdd(lhs, rhs);
		ThrowIf<IRGenerationError>(out == nullptr);
		return out;
	}

	llvm::Value* ReturnStatement::GenIR(llvm::LLVMContext& context, llvm::Module& module, llvm::IRBuilder<>& builder)
	{
		if (m_Val == nullptr)
		{
			ThrowIf<IRGenerationError>(true);
			return nullptr;
		}

		else
		{
			llvm::Value* out = builder.CreateRet(m_Val->GenIR(context, module, builder));
			ThrowIf<IRGenerationError>(out == nullptr);
			return out;
		}
	}
}

/* File 2 */

#include <Parser.h>

#include <iostream>

namespace LX
{
	void GenerateIR(FileAST& ast)
	{
		// Generates stuff //

		llvm::LLVMContext context;
		llvm::IRBuilder<> builder(context);
		
		{
			std::unique_ptr<llvm::Module> module = std::make_unique<llvm::Module>("add_ints", context);

			// Defines main function //

			llvm::FunctionType* funcType = llvm::FunctionType::get(llvm::Type::getInt32Ty(context), false);
			llvm::Function* mainFunc = llvm::Function::Create(funcType, llvm::Function::ExternalLinkage, "main", module.get());
			llvm::BasicBlock* entry = llvm::BasicBlock::Create(context, "entry", mainFunc);
			builder.SetInsertPoint(entry);

			// Loops over AST to generate IR //

			for (auto& node : ast.functions[0].body)
			{
				switch (node->m_Type)
				{
					case AST::Node::RETURN_STATEMENT:
					{
						node->GenIR(context, *module, builder);
						break;
					}

					default:
					{
						break;
					}
				}
			}

			if (entry->getTerminator() == nullptr)
			{
				builder.CreateRet(llvm::ConstantInt::get(llvm::Type::getInt32Ty(context), 0, true));
			}

			// Verification of the IR //

			if (llvm::verifyFunction(*mainFunc, &llvm::errs()) || llvm::verifyModule(*module, &llvm::errs()))
			{
				std::cerr << "Error: IR generation failed" << std::endl;
				return;
			}

			// Outputs the IR to the console //

			module->print(llvm::outs(), nullptr);

		} // <- Crashes here

		std::cout << "Finished generating IR" << std::endl;
	}
}


/* Output */

; ModuleID = 'add_ints'
source_filename = "add_ints"

define i32 @main() {
entry:
  ret i32 7
}
Finished generating IR

/* AST */

func main
    return 3 + 4