using System;
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");
			telnetPatch (module);
			module.Write ("Assembly-CSharp.dll");
			Console.WriteLine ("Done");
		}

		private static void telnetPatch (ModuleDefinition module)
		{
			TypeDefinition type = module.GetType ("NetTelnetServer");

			if (isPatched (type)) {
				return;
			}

			markTypePatched (module, type);
			patchMethod (type, ".ctor", 1, typeof(AllocsNetTelnetServer).GetMethod ("init"));
			patchMethod (type, "Disconnect", 0, typeof(AllocsNetTelnetServer).GetMethod ("Disconnect"));
			patchMethod (type, "SetConsole", 1, typeof(AllocsNetTelnetServer).GetMethod ("SetConsole"));
			patchMethod (type, "WriteToClient", 1, typeof(AllocsNetTelnetServer).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 patchMethod (TypeDefinition type, string methodName, int opCount, MethodBase targetMethod)
		{
			foreach (MethodDefinition method in type.Methods) {
				if (method.Name.Equals (methodName)) {
					Console.WriteLine ("Patching " + methodName);
					var il = method.Body.GetILProcessor ();
					var call = il.Create (OpCodes.Call, method.Module.Import (targetMethod));
					var i = 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));
				}
			}
		}

		private static bool isPatched (TypeDefinition type)
		{
			foreach (FieldDefinition fd in type.Fields) {
				if (fd.Name.Equals ("AllocsPatch")) {
					Console.WriteLine ("\"" + type.Name + "\" is already patched, skipping");
					return true;
				}
			}
			Console.WriteLine ("Patching \"" + type.Name + "\"");
			return false;
		}
	}
}