using System;
using System.Collections.Generic;
using System.Reflection;
using Mono.Cecil;
using Mono.Cecil.Cil;

namespace dtdfixer
{
	class MainClass {
		public static void Main (string[] args) {
			ModuleDefinition module = ModuleDefinition.ReadModule ("Assembly-CSharp.dll");

			TypeDefinition type = module.GetType ("GameManager");
			if (isPatched (type)) {
				Console.WriteLine ("Assembly already patched");
				return;
			}
			markTypePatched (module, type);

			mappingPatch (module);
			consoleOutputPatch (module);
			telnetPatch (module);
			connectLogPatch (module);
			playerDataPatch (module);

			module.Write ("Assembly-CSharp.patched.dll");
			Console.WriteLine ("Done");

		}

		private static void mappingPatch (ModuleDefinition module) {
			TypeDefinition type = module.GetType ("Chunk");
			addHook (type, "CalcMapColors", true, 0, true, typeof(AllocsFixes.Mods).GetMethod ("CallCalcMapColors"));
		}

		private static void consoleOutputPatch (ModuleDefinition module) {
			TypeDefinition type = module.GetType ("ConsoleSdtd");
			replaceMethod (type, "ExecuteCmdFromClient", true, 4, typeof(AllocsFixes.NetConnections.ConsoleOutputSeparator).GetMethod ("C_ExecuteCmdFromClient"));
			addHook (type, "Run", true, 0, true, typeof(AllocsFixes.NetConnections.ConsoleOutputSeparator).GetMethod ("C_Run"));
			replaceMethod (type, "SendResult", true, 1, typeof(AllocsFixes.NetConnections.ConsoleOutputSeparator).GetMethod ("C_SendResult"));
		}

		private static void playerDataPatch (ModuleDefinition module) {
			TypeDefinition type = module.GetType ("GameManager");
			addHook (type, "SavePlayerData", true, 2, true, typeof(AllocsFixes.PlayerDataStuff).GetMethod ("GM_SavePlayerData"));
			addHook (type, "Awake", true, 0, true, typeof(AllocsFixes.StateManager).GetMethod ("Awake"));
			addHook (type, "Cleanup", true, 0, false, typeof(AllocsFixes.StateManager).GetMethod ("Shutdown"));
		}

		private static void connectLogPatch (ModuleDefinition module) {
			TypeDefinition type = module.GetType ("GameManager");
			addHook (type, "RequestToSpawnPlayer", true, 4, true, typeof(AllocsFixes.AllocsLogFunctions).GetMethod ("RequestToSpawnPlayer"));
			type = module.GetType ("ConnectionManager");
			addHook (type, "DisconnectClient", true, 2, false, typeof(AllocsFixes.AllocsLogFunctions).GetMethod ("PlayerDisconnected"));
		}

		private static void telnetPatch (ModuleDefinition module) {
			TypeDefinition type = module.GetType ("NetTelnetServer");
			replaceMethod (type, ".ctor", false, 1, typeof(AllocsFixes.NetConnections.NetTelnetServer).GetMethod ("init"));
			replaceMethod (type, "Disconnect", false, 0, typeof(AllocsFixes.NetConnections.NetTelnetServer).GetMethod ("Disconnect"));
			replaceMethod (type, "SetConsole", false, 1, typeof(AllocsFixes.NetConnections.NetTelnetServer).GetMethod ("SetConsole"));
			replaceMethod (type, "WriteToClient", false, 1, typeof(AllocsFixes.NetConnections.NetTelnetServer).GetMethod ("WriteToClient"));
		}

		private static void markTypePatched (ModuleDefinition module, TypeDefinition type) {
			type.Fields.Add (new FieldDefinition ("AllocsPatch", Mono.Cecil.FieldAttributes.Private | Mono.Cecil.FieldAttributes.SpecialName, module.Import (typeof(int))));
		}

		private static void addHook (TypeDefinition type, string methodName, bool addThisRef, int opCount, bool atEnd, MethodBase targetMethod) {
			MethodDefinition method = findMethod (type, methodName, opCount);
			if (method != null) {
				var il = method.Body.GetILProcessor ();
				var call = il.Create (OpCodes.Call, method.Module.Import (targetMethod));
				if (atEnd) {
					int insBefore = method.Body.Instructions.Count;
					if (addThisRef) {
						il.Append (il.Create (OpCodes.Ldarg, 0));
					}
					for (int op = 0; op < opCount; op++) {
						il.Append (il.Create (OpCodes.Ldarg, op + 1));
					}
					il.Append (call);
					il.Remove (method.Body.Instructions [insBefore - 1]);
					il.Append (il.Create (OpCodes.Ret));
				} else {
					var i = 0;
					if (addThisRef) {
						il.InsertBefore (method.Body.Instructions [i++], il.Create (OpCodes.Ldarg, 0));
					}
					for (int op = 0; op < opCount; op++) {
						il.InsertBefore (method.Body.Instructions [i++], il.Create (OpCodes.Ldarg, op + 1));
					}
					il.InsertBefore (method.Body.Instructions [i++], call);
				}
				return;
			}
		}

		private static void replaceMethod (TypeDefinition type, string methodName, bool addThisRef, int opCount, MethodBase targetMethod) {
			MethodDefinition method = findMethod (type, methodName, opCount);
			if (method != null) {
				var il = method.Body.GetILProcessor ();
				var call = il.Create (OpCodes.Call, method.Module.Import (targetMethod));
				var i = 0;
				if (addThisRef) {
					il.InsertBefore (method.Body.Instructions [i++], il.Create (OpCodes.Ldarg, 0));
				}
				for (int op = 0; op < opCount; op++) {
					il.InsertBefore (method.Body.Instructions [i++], il.Create (OpCodes.Ldarg, op + 1));
				}
				il.InsertBefore (method.Body.Instructions [i++], call);
				il.InsertBefore (method.Body.Instructions [i++], il.Create (OpCodes.Ret));
				return;
			}
		}

		private static MethodDefinition findMethod (TypeDefinition type, string methodName, int opCount) {
			foreach (MethodDefinition method in type.Methods) {
				if (method.Name.Equals (methodName) && method.Parameters.Count == opCount) {
					return method;
				}
			}
			Console.WriteLine ("ERROR: Did not find " + type.Name + "." + methodName + "() with " + opCount + " parameters");
			return null;
		}

		private static bool isPatched (TypeDefinition type) {
			foreach (FieldDefinition fd in type.Fields) {
				if (fd.Name.Equals ("AllocsPatch")) {
					return true;
				}
			}
			return false;
		}
	}
}