// Copyright (c) 2023 bradson
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

using System.Diagnostics;
using System.Linq;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Cecil.Rocks;
using Mono.Collections.Generic;
using MonoMod.Utils;

namespace PerformanceFish.Prepatching;

public abstract class FishPrepatch : FishPrepatchBase
{
	public abstract MethodBase TargetMethodBase { get; }
	
	public virtual MethodInfo? PrefixMethodInfo
		=> TryGetMethod("PREFIX", static m => m.HasAttribute<HarmonyPrefix>());

	public virtual MethodInfo? PostfixMethodInfo
		=> TryGetMethod("POSTFIX", static m => m.HasAttribute<HarmonyPostfix>());

	public virtual void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
	{
	}

	private MethodInfo? TryGetMethod(string nameCaseInsensitive, Predicate<MethodInfo>? predicate = null)
		=> GetType().TryGetMethod(nameCaseInsensitive, predicate);

	protected internal override void ApplyPatch(ModuleDefinition module)
	{
		// ReSharper disable once NullCoalescingConditionIsAlwaysNotNullAccordingToAPIContract
		var patchMethod = TargetMethodBase
			?? ThrowHelper.ThrowInvalidOperationException<MethodBase>(
				$"Patch {GetType().FullName} has no target method.");

		if (patchMethod.IsGenericMethod && patchMethod is MethodInfo info)
			patchMethod = info.GetGenericMethodDefinition();

		var patchType = patchMethod.ReflectedType
			?? ThrowHelper.ThrowInvalidOperationException<Type>(
				$"Method {patchMethod.FullDescription()} has no declaring type.");

		if (patchMethod.ReflectedType != patchMethod.DeclaringType)
		{
			ThrowHelper.ThrowInvalidOperationException($"Tried to apply patch on method {
				patchMethod.Name} for type {patchMethod.ReflectedType}, but it is only declared in {
					patchMethod.DeclaringType}.");
		}

		if (patchType.IsGenericType)
		{
			patchType = patchType.GetGenericTypeDefinition();
			patchMethod = patchMethod is MethodInfo
				? AccessTools.DeclaredMethod(patchType, patchMethod.Name, patchMethod.GetParameters().GetTypes())
				: AccessTools.Constructor(patchType, patchMethod.GetParameters().GetTypes(), patchMethod.IsStatic);
		}
		
		DebugLog.Message($"Applying prepatch {GetType().FullName} on method {
			patchMethod.FullDescription()}");

		var targetType = module.GetTypeDefinition(patchType) ?? ThrowForFailureToFindType(patchType);
		var targetMethod = targetType.GetMethod(patchMethod.Name, patchMethod.GetParameters().GetTypes())
			?? ThrowForFailureToFindMethod(targetType, patchMethod);
		
		// PrepatchManager.AddSecurityAttributes(targetMethod.SecurityDeclarations, module, targetMethod.FullName);

		var targetMethodBody = targetMethod.Body;
		var ilProcessor = targetMethodBody.GetILProcessor();
		
		try
		{
			Transpiler(ilProcessor, module);
		}
		catch (Exception e)
		{
			Log.Error($"Exception while applying transpiler for {GetType().FullName}\n{e}\n{new StackTrace(true)}");
		}

		var patchVariables = new PatchVariables();

		if (PrefixMethodInfo is { } prefixMethod)
		{
			try
			{
				ApplyPrefix(prefixMethod, ilProcessor, module, ref patchVariables);
			}
			catch (Exception e)
			{
				Log.Error($"Exception while applying prefix for {GetType().FullName}\n{e}\n{new StackTrace(true)}");
			}
		}

		if (PostfixMethodInfo is { } postfixMethod)
		{
			try
			{
				ApplyPostfix(postfixMethod, ilProcessor, module, ref patchVariables);
			}
			catch (Exception e)
			{
				Log.Error($"Exception while applying postfix for {GetType().FullName}\n{e}\n{new StackTrace(true)}");
			}
		}

		targetMethodBody.SimplifyMacros();
		targetMethodBody.OptimizeMacros();
	}

	private static void ApplyPrefix(MethodInfo prefixMethodInfo, ILProcessor ilProcessor, ModuleDefinition module,
		ref PatchVariables patchVariables)
	{
		if (prefixMethodInfo.ReturnType != typeof(bool) && prefixMethodInfo.ReturnType != typeof(void))
		{
			ThrowHelper.ThrowArgumentException($"Prefix has invalid return type: {
				prefixMethodInfo.ReturnType.FullName}. Must be void or bool");
		}

		var instructions = ilProcessor.instructions;
		var targetMethod = ilProcessor.GetMethod();

		var genericParameters = GetGenericParameters(targetMethod);
		var prefixParameters = prefixMethodInfo.GetParameters();
		
		var resultParameter = prefixParameters.TryGetNamed("__result");
		VerifyResultParameter(module, resultParameter, targetMethod, true);

		var stateParameter = prefixParameters.TryGetNamed("__state");
		VerifyStateParameter(stateParameter);

		// var skipOriginalFlag = ilProcessor.DeclareLocal(typeof(bool));
		var resultVariable = GetResultVariable(ilProcessor, ref patchVariables, resultParameter, targetMethod);
		var stateVariable = GetStateVariable(ilProcessor, ref patchVariables, stateParameter);
		var startOfOriginal = instructions[0];

		var index = 0;
		foreach (var parameter in prefixParameters)
		{
			ilProcessor.InsertAt(index++, GetInstructionForParameter(module, parameter, resultParameter, resultVariable,
				stateParameter, stateVariable, targetMethod));
		}
		
		if (genericParameters != null)
		{
			ilProcessor.InsertAt(index++, OpCodes.Call,
				module.ImportMethodWithGenericArguments(prefixMethodInfo, genericParameters));
		}
		else
		{
			ilProcessor.InsertAt(index++, OpCodes.Call, prefixMethodInfo);
		}

		if (prefixMethodInfo.ReturnType == typeof(bool))
		{
			ilProcessor.InsertAt(index++, (OpCodes.Brtrue, startOfOriginal));
			
			if (resultVariable != null)
			{
				ilProcessor.InsertAt(index++, (OpCodes.Ldloc, resultVariable));
			}
			else if (targetMethod.ReturnType.FullName != module.TypeSystem.Void.FullName)
			{
				Log.Error($"Prefix {
					prefixMethodInfo.FullDescription()} with return type of bool must assign to a __result parameter");
			}

			ilProcessor.InsertAt(index, OpCodes.Ret);
		}
	}
	
	private static void ApplyPostfix(MethodInfo postfixMethodInfo, ILProcessor ilProcessor, ModuleDefinition module,
		ref PatchVariables patchVariables)
	{
		var targetMethod = ilProcessor.GetMethod();
		
		if (postfixMethodInfo.ReturnType != typeof(void)
			&& !postfixMethodInfo.ReturnType.Is(targetMethod.ReturnType))
		{
			ThrowHelper.ThrowArgumentException($"Postfix has invalid return type: {
				postfixMethodInfo.ReturnType.FullName}. Must be void or {targetMethod.ReturnType}");
		}

		var instructions = ilProcessor.instructions;
		var genericParameters = GetGenericParameters(targetMethod);
		var postfixParameters = postfixMethodInfo.GetParameters();
		
		var resultParameter = postfixParameters.TryGetNamed("__result");
		VerifyResultParameter(module, resultParameter, targetMethod, false);

		var stateParameter = postfixParameters.TryGetNamed("__state");
		// Verify? It would make sense to match against prefix state

		var hasReturnValue = targetMethod.ReturnType.FullName != module.TypeSystem.Void.FullName;
		var resultVariable = hasReturnValue
			? patchVariables.Result ??= ilProcessor.DeclareLocal(targetMethod.ReturnType)
			: GetResultVariable(ilProcessor, ref patchVariables, resultParameter, targetMethod);
		var stateVariable = GetStateVariable(ilProcessor, ref patchVariables, stateParameter);
		var retInstructions = instructions.Where(static instruction => instruction.OpCode == OpCodes.Ret).ToArray();
		if (retInstructions.Length == 0)
			return;

		var postfixInstructions = new Collection<object>();
		
		if (hasReturnValue)
		{
			postfixInstructions.Add((OpCodes.Stloc, resultVariable!));
			postfixInstructions.Add(OpCodes.Nop);
		}
		else
		{
			postfixInstructions.Add(OpCodes.Nop);
		}

		foreach (var parameter in postfixParameters)
		{
			postfixInstructions.Add(GetInstructionForParameter(module, parameter, resultParameter, resultVariable,
				stateParameter, stateVariable, targetMethod));
		}
		
		if (genericParameters != null)
		{
			postfixInstructions.Add((OpCodes.Call,
				module.ImportMethodWithGenericArguments(postfixMethodInfo, genericParameters)));
		}
		else
		{
			postfixInstructions.Add((OpCodes.Call, postfixMethodInfo));
		}

		if (postfixMethodInfo.ReturnType != typeof(void))
			postfixInstructions.Add((OpCodes.Stloc, resultVariable!));
		
		if (hasReturnValue)
			postfixInstructions.Add((OpCodes.Ldloc, resultVariable!));

		var lastRetInstruction = retInstructions[^1];
		var lastRetInstructionIndex = instructions.IndexOf(lastRetInstruction);
		ilProcessor.InsertRange(lastRetInstructionIndex, postfixInstructions);
		var captureStart = instructions[lastRetInstructionIndex];
		var postfixStart = instructions[lastRetInstructionIndex + (hasReturnValue ? 1 : 0)];

		RetargetExceptionHandlerEnds(targetMethod.Body, lastRetInstruction, captureStart);
		var retInstructionSet = retInstructions.ToHashSet();
		foreach (var instruction in instructions)
		{
			if (instruction.Operand is Instruction branchTarget && retInstructionSet.Contains(branchTarget))
				instruction.Operand = captureStart;
		}

		for (var i = 0; i < retInstructions.Length - 1; i++)
		{
			var instruction = retInstructions[i];
			if (hasReturnValue)
				ilProcessor.InsertBefore(instruction, Instruction.Create(OpCodes.Stloc, resultVariable!));

			instruction.OpCode = IsInProtectedRegion(targetMethod.Body, instruction) ? OpCodes.Leave : OpCodes.Br;
			instruction.Operand = postfixStart;
		}
	}

	private static void RetargetExceptionHandlerEnds(Mono.Cecil.Cil.MethodBody methodBody, Instruction oldBoundary,
		Instruction newBoundary)
	{
		foreach (var exceptionHandler in methodBody.ExceptionHandlers)
		{
			if (exceptionHandler.TryEnd == oldBoundary)
				exceptionHandler.TryEnd = newBoundary;
			if (exceptionHandler.HandlerEnd == oldBoundary)
				exceptionHandler.HandlerEnd = newBoundary;
		}
	}

	private static bool IsInProtectedRegion(Mono.Cecil.Cil.MethodBody methodBody, Instruction instruction)
	{
		var instructions = methodBody.Instructions;
		var instructionIndex = instructions.IndexOf(instruction);
		
		foreach (var exceptionHandler in methodBody.ExceptionHandlers)
		{
			if (IsWithinRange(instructions, instructionIndex, exceptionHandler.TryStart, exceptionHandler.TryEnd)
				|| IsWithinRange(instructions, instructionIndex, exceptionHandler.HandlerStart,
					exceptionHandler.HandlerEnd)
				|| exceptionHandler.FilterStart != null
				&& IsWithinRange(instructions, instructionIndex, exceptionHandler.FilterStart,
					exceptionHandler.HandlerStart))
			{
				return true;
			}
		}

		return false;
	}

	private static bool IsWithinRange(Mono.Collections.Generic.Collection<Instruction> instructions,
		int instructionIndex, Instruction? rangeStart, Instruction? rangeEnd)
	{
		if (rangeStart == null)
			return false;

		var rangeStartIndex = instructions.IndexOf(rangeStart);
		if (rangeStartIndex < 0 || instructionIndex < rangeStartIndex)
			return false;

		if (rangeEnd == null)
			return true;

		var rangeEndIndex = instructions.IndexOf(rangeEnd);
		return rangeEndIndex < 0 || instructionIndex < rangeEndIndex;
	}

	private static VariableDefinition? GetStateVariable(ILProcessor ilProcessor, ref PatchVariables patchVariables,
		ParameterInfo? stateParameter)
		=> stateParameter != null
			? patchVariables.State ??= ilProcessor.DeclareLocal(GetUnderlyingType(stateParameter.ParameterType))
			: null;

	private static VariableDefinition? GetResultVariable(ILProcessor ilProcessor, ref PatchVariables patchVariables,
		ParameterInfo? resultParameter, MethodDefinition targetMethod)
		=> resultParameter != null
			? patchVariables.Result ??= ilProcessor.DeclareLocal(targetMethod.ReturnType)
			: null;

	private static void VerifyStateParameter(ParameterInfo? stateParameter)
	{
		if (stateParameter is { ParameterType.IsByRef: false })
			ThrowHelper.ThrowArgumentException("__state parameter must be declared with ref or out keyword.");
	}

	private static ITuple GetInstructionForParameter(ModuleDefinition module, ParameterInfo parameter,
		ParameterInfo? resultParameter, VariableDefinition? resultVariable, ParameterInfo? stateParameter,
		VariableDefinition? stateVariable, MethodDefinition targetMethod)
		=> parameter == resultParameter
			? (parameter.ParameterType.IsByRef ? OpCodes.Ldloca : OpCodes.Ldloc, resultVariable)
			: parameter == stateParameter
				? (parameter.ParameterType.IsByRef ? OpCodes.Ldloca : OpCodes.Ldloc, stateVariable)
				: GetMatchingTargetParameter(targetMethod, parameter, module);

	private static void VerifyResultParameter(ModuleDefinition module, ParameterInfo? resultParameter,
		MethodDefinition targetMethod, bool requiresByRef)
	{
		if (resultParameter is null)
			return;

		if (targetMethod.ReturnType.IsGeneric() ? !GetUnderlyingType(resultParameter.ParameterType).IsGeneric()
			: !GetUnderlyingType(resultParameter.ParameterType).Is(targetMethod.ReturnType))
		{
			// var firstType = targetMethod.ReturnType;
			// var secondType = GetUnderlyingType(resultParameter.ParameterType);
			// Log.Message($"Type FullName: {firstType.FullName
			// }: ContainsGenericParameter: {firstType.ContainsGenericParameter
			// }, HasGenericParameters :{firstType.HasGenericParameters
			// }, IsGenericInstance: {firstType.IsGenericInstance
			// }, Type FullName: {secondType.FullName
			// } ContainsGenericParameters: {secondType.ContainsGenericParameters
			// }, IsGenericType {secondType.IsGenericType
			// }, IsConstructedGenericType: {secondType.IsConstructedGenericType}");
			
			ThrowHelper.ThrowArgumentException($"__result parameter has invalid type: {
				resultParameter.ParameterType.FullName}. Must match the target's return type: {
					targetMethod.ReturnType.FullName}");
		}
		if (requiresByRef && !resultParameter.ParameterType.IsByRef)
		{
			ThrowHelper.ThrowArgumentException(
				"__result parameter must be declared with ref or out keyword.");
		}
	}

	private static Collection<GenericParameter>? GetGenericParameters(MethodDefinition targetMethod)
		=> targetMethod.HasGenericParameters ? targetMethod.GenericParameters
			: targetMethod.DeclaringType.HasGenericParameters ? targetMethod.DeclaringType.GenericParameters
			: null;

	private static Type GetUnderlyingType(Type type) => type.IsByRef ? type.GetElementType()! : type;

	private static ITuple GetMatchingTargetParameter(MethodDefinition targetMethod,
		ParameterInfo parameter, ModuleDefinition module)
	{
		var targetParameters = targetMethod.Parameters;
		var targetParameter = parameter.Name == "__instance"
			? targetMethod.Body.ThisParameter
			: targetParameters.TryGetNamed(parameter.Name);
		if (targetParameter is null)
		{
			return ThrowHelper.ThrowArgumentException<ITuple>($"No parameter named {
				parameter.Name} found in target method");
		}

		var targetParameterType = targetParameter.ParameterType;
		var patchParameterType = parameter.ParameterType;

		return GetUnderlyingType(targetParameter).Name != GetUnderlyingType(patchParameterType).Name
			? ThrowHelper.ThrowArgumentException<ITuple>($"Parameter has invalid type: {
				patchParameterType.FullName}. Must match the target method parameter's type: {
					targetParameterType.FullName}")
			: (patchParameterType.IsByRef && !targetParameter.ParameterType.IsByReference
				? OpCodes.Ldarga
				: OpCodes.Ldarg, targetParameter);
	}

	private static TypeReference GetUnderlyingType(ParameterInfo prefixParameter, ModuleDefinition module)
		=> module.ImportReference(GetUnderlyingType(prefixParameter.ParameterType));

	private static TypeReference GetUnderlyingType(ParameterDefinition prefixParameter)
		=> prefixParameter.ParameterType.IsByReference ? prefixParameter.ParameterType.GetElementType()!
			: prefixParameter.ParameterType;

	private static MethodDefinition ThrowForFailureToFindMethod(TypeDefinition? type, MethodBase patchMethod)
		=> ThrowHelper.ThrowInvalidOperationException<MethodDefinition>(
			$"Failed to find method with type {type} and name {patchMethod.Name}");

	private record struct PatchVariables
	{
		public VariableDefinition? Result, State;
	}
}