// Copyright Epic Games, Inc. All Rights Reserved.

using EpicGames.Core;
using Microsoft.Extensions.Logging;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace UnrealBuildTool
{
	namespace VerseVMBytecode
	{
		internal enum CppType
		{
			LabelOffset,
		}

		internal enum Role
		{
			Use,
			UnifyDef,
			ClobberDef,
		}

		internal enum Arity
		{
			Fixed,
			Variadic,
		}

		static class Extensions
		{
			public static string ToCpp(this Role TheRole)
			{
				switch (TheRole)
				{
					case Role.Use:
						return "EOperandRole::Use";
					case Role.UnifyDef:
						return "EOperandRole::UnifyDef";
					case Role.ClobberDef:
						return "EOperandRole::ClobberDef";
				}
				return "#error";
			}

			public static bool IsDef(this Role TheRole)
			{
				return TheRole == Role.UnifyDef || TheRole == Role.ClobberDef;
			}

			public static string DefCppType(this Role TheRole)
			{
				return TheRole.IsDef() ? "FRegisterIndex" : "FValueOperand";
			}

			public static string ToCpp(this CppType Type)
			{
				switch (Type)
				{
					case CppType.LabelOffset:
						return "FLabelOffset";
				}
				return "#error";
			}

			public static string ToCpp(this bool Bool)
			{
				return Bool ? "true" : "false";
			}
		}

		internal class Argument
		{
			public string Name;
			public Role Role;
			public Arity Arity;

			public Argument(in string _Name, in Role _Role, in Arity _Arity)
			{
				Name = _Name;
				Role = _Role;
				Arity = _Arity;
			}

			public Argument(in string _Name, in Role _Role) : this(_Name, _Role, Arity.Fixed)
			{ }
		}

		internal class Constant
		{
			public string Name;
			public CppType Type;

			public Constant(string _Name, CppType _Type)
			{
				Name = _Name;
				Type = _Type;
			}

			public bool IsJump()
			{
				return Type == CppType.LabelOffset;
			}
		}

		internal class Instruction
		{
			public string Name;
			public List<Argument> Args = new List<Argument>();
			public List<Constant> Consts = new List<Constant>();
			public bool _CapturesEffectToken = false;
			public bool _CreatesNewReturnEffectToken = false;
			public bool _Suspends = false;
			public bool _Jumps = false;

			public Instruction(string _Name)
			{
				Name = _Name;
			}

			public string CppName => $"FOp{Name}";

			public string CppCapturesName => $"F{Name}SuspensionCaptures";

			public Instruction Arg(in string Name, in Role InRole, in Arity InArity)
			{
				Args.Add(new Argument(Name, InRole, InArity));
				return this;
			}

			public Instruction Arg(in string Name, in Role InRole)
			{
				return Arg(Name, InRole, Arity.Fixed);
			}

			public Instruction Const(string Name, CppType Type)
			{
				Consts.Add(new Constant(Name, Type));
				return this;
			}

			public Instruction Jump(string Name)
			{
				_Jumps = true;
				return Const(Name, CppType.LabelOffset);
			}

			public Instruction CapturesEffectToken()
			{
				_CapturesEffectToken = true;
				return this;
			}

			public Instruction CreatesNewReturnEffectToken()
			{
				_CreatesNewReturnEffectToken = true;
				return this;
			}

			public Instruction Suspends()
			{
				_Suspends = true;
				return this;
			}
		}
	}
}

namespace UnrealBuildTool
{
	using VerseVMBytecode;

	/// <summary>
	/// Generates bytecode and bytecode helpers for the VerseVM.
	/// </summary>
	public class VerseVMBytecodeGenerator
	{
		readonly List<Instruction> Instructions = new List<Instruction>();

		Instruction Inst(string Name)
		{
			Instruction I = new Instruction(Name);
			Instructions.Add(I);
			return I;
		}

		static string Preamble()
		{
			StringBuilder S = new StringBuilder();
			S.Append("// Copyright Epic Games, Inc. All Rights Reserved.\n\n");
			S.Append("// WARNING: This code is autogenerated by VerseVMBytecodeGenerator.cs. Do not edit directly\n\n");
			S.Append("#pragma once\n\n");
			return S.ToString();
		}

		string EmitBytecodeMacroList()
		{
			StringBuilder S = new StringBuilder();
			S.Append(Preamble());

			S.Append("// IWYU pragma: private, include \"VVMBytecodeOps.h\"\n\n");

			S.Append("#define VERSE_ENUM_OPS(v) \\\n");
			foreach (Instruction Inst in Instructions)
			{
				S.Append($"    v({Inst.Name}) \\\n");
			}
			S.Append("\n");

			return S.ToString();
		}

		string EmitReflectionMethods(in Instruction Inst, in bool bIsSuspensionCapture)
		{
			StringBuilder S = new StringBuilder();
			S.Append("    template <typename FunctionType>\n");
			S.Append("    void ForEachOperand(FunctionType Function) const\n");
			S.Append("    {\n");
			foreach (Argument Arg in Inst.Args)
			{
				string ArgString = bIsSuspensionCapture ? $"{Arg.Name}.Get()" : $"{Arg.Name}";
				string OperandString = bIsSuspensionCapture ? "Operand.Get()" : "Operand";
				if (Arg.Arity == Arity.Variadic)
				{
					S.Append($"        for (const auto& Operand : {Arg.Name})\n");
					S.Append("        {\n");
					S.Append($"            Function({Arg.Role.ToCpp()}, {OperandString});\n");
					S.Append("        }\n");
				}
				else
				{
					S.Append($"        Function({Arg.Role.ToCpp()}, {ArgString});\n");
				}
			}
			S.Append("    }\n\n");
			S.Append("    template <typename FunctionType>\n");
			S.Append("    void ForEachOperandWithName(FunctionType Function) const\n");
			S.Append("    {\n");
			foreach (Argument Arg in Inst.Args)
			{
				string ArgString = bIsSuspensionCapture ? $"{Arg.Name}.Get()" : $"{Arg.Name}";
				string OperandString = bIsSuspensionCapture ? "Operand.Get()" : "Operand";
				if (Arg.Arity == Arity.Variadic)
				{
					S.Append($"        for (int32 Index = 0; Index < {Arg.Name}.Num(); ++Index)\n");
					S.Append("        {\n");
					S.Append($"            auto& Operand = {Arg.Name}[Index];\n");
					S.Append($"            const FString ArgName = FString::Format(TEXT(\"{{0}}{{1}}\"), {{\"{Arg.Name}\", Index}});\n");
					S.Append($"            Function({Arg.Role.ToCpp()}, {OperandString}, TCHAR_TO_ANSI(*ArgName));\n");
					S.Append("        }\n");
				}
				else
				{
					S.Append($"        Function({Arg.Role.ToCpp()}, {ArgString}, \"{Arg.Name}\");\n");
				}
			}
			S.Append("    }\n\n");
			return S.ToString();
		}

		String EmitBytecodeAndCaptureDefs()
		{
			StringBuilder S = new StringBuilder();
			S.Append(Preamble());

			S.Append("// IWYU pragma: private, include \"VVMBytecodesAndCaptures.h\"\n\n");

			S.Append("namespace Verse {\n");

			// Emit bytecode structs.
			foreach (Instruction Inst in Instructions)
			{
				// Bytecode structs are 1-byte aligned. If we expand them to have intermediates that can
				// be written to dynamically (a la an IC), and are also scanned by a concurrent GC, those
				// fields will need to have natural alignment to avoid the GC seeing torn values. 
				S.Append($"#pragma pack(push, 1)\n");

				S.Append($"struct {Inst.CppName} : public FOp");
				S.Append("\n{\n");

				// Define fields
				foreach (Argument Arg in Inst.Args)
				{
					switch (Arg.Arity)
					{
						case Arity.Variadic:
							// NOTE: (yiliang.siew) This could be raised to a location such as in `VProgram` when that
							// exists and each opcode store only the index + size to index into that array, so that we don't
							// need to store a separate `TArray` of operand values per-opcode struct.
							S.Append($"    TArray<{Arg.Role.DefCppType()}> {Arg.Name};  // Variadic argument.\n");
							break;
						case Arity.Fixed:
							S.Append($"    {Arg.Role.DefCppType()} {Arg.Name};\n");
							break;
						default:
							throw new ArgumentException("Invalid argument type!");
					}
				}
				foreach (Constant Const in Inst.Consts)
				{
					S.Append($"    {Const.Type.ToCpp()} {Const.Name};\n");
				}
				S.Append($"#pragma pack(pop)\n\n");

				S.Append($"    static constexpr EOpcode StaticOpcode = EOpcode::{Inst.Name};\n");
				S.Append($"    static constexpr bool bHasJumps = {Inst._Jumps.ToCpp()};\n\n");

				// Constructor
				S.Append($"    {Inst.CppName}(");
				string ArgumentsList = string.Join(", ", Inst.Args.Select(Arg => Arg.Arity == Arity.Variadic ? $"TArray<{Arg.Role.DefCppType()}>&& {Arg.Name}" : $"const {Arg.Role.DefCppType()} {Arg.Name}"));
				S.Append(ArgumentsList);
				string ConstantsList = string.Join(", ", Inst.Consts.Select(Const => $"{Const.Type.ToCpp()} {Const.Name}"));
				S.Append(ConstantsList);
				S.Append(")\n");
				S.Append("        : FOp(StaticOpcode)\n");
				foreach (Argument Arg in Inst.Args)
				{
					if (Arg.Arity == Arity.Variadic)
					{
						S.Append($"        , {Arg.Name}(MoveTemp({Arg.Name}))\n");
					}
					else
					{
						S.Append($"        , {Arg.Name}({Arg.Name})\n");
					}
				}
				foreach (Constant Const in Inst.Consts)
				{
					S.Append($"        , {Const.Name}({Const.Name})\n");
				}
				S.Append("    {}\n\n");

				// Reflection methods 
				S.Append(EmitReflectionMethods(Inst, false));
				S.Append("    template <typename FunctionType>\n");
				S.Append("    FORCEINLINE void ForEachJump(FunctionType Function) const\n");
				S.Append("    {\n");
				foreach (Constant Const in Inst.Consts.Where(C => C.IsJump()))
				{
					S.Append($"        Function({Const.Name});\n");
				}
				S.Append("    }\n");

				S.Append("    template <typename FunctionType>\n");
				S.Append("    FORCEINLINE void ForEachJumpWithName(FunctionType Function) const\n");
				S.Append("    {\n");
				foreach (Constant Const in Inst.Consts.Where(C => C.IsJump()))
				{
					S.Append($"        Function({Const.Name}, \"{Const.Name}\");\n");
				}
				S.Append("    }\n");

				S.Append("\n};\n\n");
			}

			// Emit captures structs.
			foreach (Instruction Inst in Instructions.Where(I => I._Suspends))
			{
				string Name = Inst.CppCapturesName;
				S.Append($"struct {Name}");
				S.Append("\n{\n");

				// Generate the fields.
				foreach (Argument Arg in Inst.Args)
				{
					if (Arg.Arity == Arity.Variadic)
					{
						S.Append($"    TArray<TWriteBarrier<VValue>> {Arg.Name};  // Captured variadic arguments.\n");
					}
					else
					{
						S.Append($"    TWriteBarrier<VValue> {Arg.Name};\n");
					}
				}
				if (Inst._CapturesEffectToken)
				{
					S.Append($"    TWriteBarrier<VValue> EffectToken;\n");
				}
				if (Inst._CreatesNewReturnEffectToken)
				{
					S.Append($"    TWriteBarrier<VValue> ReturnEffectToken;\n");
				}

				S.Append("\n");

				// Generate the constructor.
				{
					S.Append($"    {Name}(FAccessContext Context");
					foreach (Argument Arg in Inst.Args)
					{
						if (Arg.Arity == Arity.Variadic)
						{
							// Avoid variable shadowing.
							S.Append($", TArray<TWriteBarrier<VValue>>&& In{Arg.Name}");
						}
						else
						{
							S.Append($", VValue {Arg.Name}");
						}
					}
					if (Inst._CapturesEffectToken)
					{
						S.Append(", VValue EffectToken");
					}
					if (Inst._CreatesNewReturnEffectToken)
					{
						S.Append(", VValue ReturnEffectToken");
					}
					S.Append(")\n");
					string Prefix = ":";
					foreach (Argument Arg in Inst.Args)
					{
						if (Arg.Arity == Arity.Variadic)
						{
							S.Append($"        {Prefix} {Arg.Name}(MoveTemp(In{Arg.Name}))\n");
						}
						else
						{
							S.Append($"        {Prefix} {Arg.Name}(Context, {Arg.Name})\n");
						}
						Prefix = ",";
					}
					if (Inst._CapturesEffectToken)
					{
						S.Append($"        {Prefix} EffectToken(Context, EffectToken)\n");
						Prefix = ", ";
					}
					if (Inst._CreatesNewReturnEffectToken)
					{
						S.Append($"        {Prefix} ReturnEffectToken(Context, ReturnEffectToken)\n");
						Prefix = ", ";
					}
					S.Append("    {}\n");
				}

				S.Append("\n");

				// Generate the copy constructor.
				{
					S.Append($"    {Name}(FAccessContext Context, const {Name}& Other)\n");
					string Prefix = ":";
					foreach (Argument Arg in Inst.Args)
					{
						if (Arg.Arity == Arity.Variadic)
						{
							S.Append($"        {Prefix} {Arg.Name}(Other.{Arg.Name})\n");
						}
						else
						{
							S.Append($"        {Prefix} {Arg.Name}(Context, Other.{Arg.Name}.Get())\n");
						}
						Prefix = ",";
					}
					if (Inst._CapturesEffectToken)
					{
						S.Append($"        {Prefix} EffectToken(Context, Other.EffectToken.Get())\n");
						Prefix = ", ";
					}
					if (Inst._CreatesNewReturnEffectToken)
					{
						S.Append($"        {Prefix} ReturnEffectToken(Context, Other.ReturnEffectToken.Get())\n");
						Prefix = ", ";
					}
					S.Append("    {\n    }\n");
				}

				S.Append("\n");
				S.Append(EmitReflectionMethods(Inst, true));
				S.Append("};\n\n");
			}

			S.Append("} // namespace Verse\n");

			return S.ToString();
		}

		string EmitMakeCapturesFunctions()
		{
			StringBuilder S = new StringBuilder();
			S.Append(Preamble());

			foreach (Instruction Inst in Instructions.Where(I => I._Suspends))
			{
				S.Append($"FORCEINLINE {Inst.CppCapturesName} MakeCaptures(const {Inst.CppName}& Op)\n{{\n");

				if (Inst._CapturesEffectToken)
				{
					S.Append("    const VValue IncomingEffectToken = EffectToken.Get(Context);\n");
				}
				if (Inst._CreatesNewReturnEffectToken)
				{
					S.Append("    const VValue ReturnEffectToken = VValue::Placeholder(VPlaceholder::New(Context, 0));\n");
					S.Append("    EffectToken.Set(Context, ReturnEffectToken);\n");
				}
				foreach (Argument Arg in Inst.Args)
				{
					if (Arg.Arity == Arity.Variadic)
					{
						S.Append($"    TArray<TWriteBarrier<VValue>> Array{Arg.Name};\n");
						S.Append($"    for (auto& CurrentValue : Op.{Arg.Name})\n");
						S.Append("    {\n");
						S.Append($"        Array{Arg.Name}.Add({{Context, GetOperand(CurrentValue)}});\n");
						S.Append("    }\n");
					}
				}
				S.Append($"    return {Inst.CppCapturesName}(Context");
				foreach (Argument Arg in Inst.Args)
				{
					if (Arg.Arity == Arity.Variadic)
					{
						S.Append($", MoveTemp(Array{Arg.Name})");
					}
					else
					{
						S.Append($", GetOperand(Op.{Arg.Name})");
					}
				}
				if (Inst._CapturesEffectToken)
				{
					S.Append(", IncomingEffectToken");
				}
				if (Inst._CreatesNewReturnEffectToken)
				{
					S.Append(", ReturnEffectToken");
				}
				S.Append(");\n");

				S.Append("}\n\n");
			}

			return S.ToString();
		}

		string EmitCaptureSwitch()
		{
			StringBuilder S = new StringBuilder();
			S.Append(Preamble());

			S.Append("// IWYU pragma: private, include \"VVMCaptureSwitch.h\"\n\n");

			S.Append("namespace Verse {\n");
			S.Append("template <typename TFunc>\n");
			S.Append("void VBytecodeSuspension::CaptureSwitch(const TFunc& Func)\n");
			S.Append("{\n");
			S.Append("    switch (PC->Opcode)\n");
			S.Append("    {\n");
			foreach (Instruction Inst in Instructions.Where(I => I._Suspends))
			{
				S.Append($"        case EOpcode::{Inst.Name}:\n");
				S.Append("        {\n");
				S.Append($"            Func(GetCaptures<{Inst.CppCapturesName}>());\n");
				S.Append("            break;\n");
				S.Append("        }\n");
			}
			S.Append("        default:\n");
			S.Append("        {\n");
			S.Append("            V_DIE(\"Opcode doesn't have a captures\");\n");
			S.Append("            break;\n");
			S.Append("        }\n");
			S.Append("    }\n");
			S.Append("}\n");
			S.Append("} // namespace Verse\n");

			return S.ToString();
		}

		VerseVMBytecodeGenerator(ILogger Logger, DirectoryReference GenDirectory)
		{
			DefineOps();

			Action<string, Func<string>> GenFile = (FileName, Method) =>
			{
				FileReference File = FileReference.Combine(GenDirectory, FileName);
				bool bWritten = FileReference.WriteAllTextIfDifferent(File, Method());
				Logger.LogDebug($"\tWriting out generated header file. Changed:{bWritten} Path:'{File}'");
			};

			GenFile("VVMBytecodeOps.gen.h", EmitBytecodeMacroList);
			GenFile("VVMBytecodesAndCaptures.gen.h", EmitBytecodeAndCaptureDefs);
			GenFile("VVMMakeCapturesFuncs.gen.h", EmitMakeCapturesFunctions);
			GenFile("VVMCaptureSwitch.gen.h", EmitCaptureSwitch);
		}

		/// <summary>
		/// Entrypoint to generate the bytecode. Generated code will go in Directory.
		/// </summary>
		public static void Generate(ILogger Logger, DirectoryReference Directory)
		{
			Logger.LogDebug($"VerseVMBytecodeGenerator.Generate, generating CPP headers in: '{Directory}'");
			new VerseVMBytecodeGenerator(Logger, Directory);
		}

		void DefineOps()
		{
			string[] BinOps =
			{
				"Add", "Sub", "Mul", "Div", "Mod"
			};
			foreach (string Op in BinOps)
			{
				Inst(Op)
					.Arg("Dest", Role.UnifyDef)
					.Arg("LeftSource", Role.Use)
					.Arg("RightSource", Role.Use)
					.Suspends();
			}

			string[] UnaryOps =
			{
				"Neg",
				"Query"
			};
			foreach (string Op in UnaryOps)
			{
				Inst(Op)
					.Arg("Dest", Role.UnifyDef)
					.Arg("Source", Role.Use)
					.Suspends();
			}

			Inst("Err");

			Inst("Move")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Source", Role.Use);

			Inst("Reset")
				.Arg("Dest", Role.ClobberDef);

			Inst("Jump")
				.Jump("JumpOffset");

			// These labels are needed for lenient execution.
			// On success, EndFailureContext falls through.
			// On failure, it jumps to OnFailure.
			// When there are still unresolved suspensions in this failure context,
			// we jump to "Done" to continue lenient execution.
			Inst("BeginFailureContext")
				.Jump("OnFailure");
			Inst("EndFailureContext")
				.Jump("Done");

			Inst("Call")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Callee", Role.Use)
				.Arg("Arguments", Role.Use, Arity.Variadic)
				.CapturesEffectToken()
				.CreatesNewReturnEffectToken()
				.Suspends();

			Inst("Return")
				.Arg("Value", Role.Use);

			Inst("NewVar")
				.Arg("Dest", Role.UnifyDef);
			Inst("VarGet")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Var", Role.Use)
				.CapturesEffectToken()
				.Suspends();
			Inst("VarSet")
				.Arg("Var", Role.Use)
				.Arg("Value", Role.Use)
				.CapturesEffectToken()
				.Suspends();

			Inst("NewTuple")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Values", Role.Use, Arity.Variadic);

			Inst("Length")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Container", Role.Use)
				.Suspends();

			Inst("IndexSet")
				.Arg("Container", Role.Use)
				.Arg("Index", Role.Use)
				.Arg("ValueToSet", Role.Use)
				.CapturesEffectToken()
				.Suspends();

			Inst("NewArrayWithCapacity")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Size", Role.Use)
				.Suspends();

			Inst("NewArray")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Values", Role.Use, Arity.Variadic);

			Inst("ArrayAdd")
				.Arg("Container", Role.Use)
				.Arg("ValueToAdd", Role.Use)
				.Suspends();

			Inst("NewOption")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Value", Role.Use)
				.Suspends();

			Inst("NewMap")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Keys", Role.Use, Arity.Variadic)
				.Arg("Values", Role.Use, Arity.Variadic)
				.Suspends();

			Inst("MapKey")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Map", Role.Use)
				.Arg("Index", Role.Use)
				.Suspends();

			Inst("MapValue")
				.Arg("Dest", Role.UnifyDef)
				.Arg("Map", Role.Use)
				.Arg("Index", Role.Use)
				.Suspends();

			string[] ComparisonOps =
			{
				"Neq", "Lt", "Lte", "Gt", "Gte"
			};
			foreach (string Op in ComparisonOps)
			{
				Inst(Op)
					.Arg("Dest", Role.UnifyDef)
					.Arg("LeftSource", Role.Use)
					.Arg("RightSource", Role.Use)
					.Suspends();
			}
		}
	}
}