// Copyright (c) 2026 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 Mono.Cecil;
using Mono.Cecil.Cil;
using PerformanceFish.Prepatching;

namespace PerformanceFish;

public sealed class CompTickOptimization : ClassWithFishPrepatches
{
#if V1_6
	public struct ThingTickCache
	{
		public List<ThingComp>? Source;
		public int Version;
		public bool AllTrivial;
	}
#endif

	private static readonly MethodInfo _thingCompTickMethod
		= AccessTools.DeclaredMethod(typeof(ThingComp), nameof(ThingComp.CompTick));

	private static readonly MethodInfo _thingCompTickIntervalMethod
		= AccessTools.DeclaredMethod(typeof(ThingComp), nameof(ThingComp.CompTickInterval), [typeof(int)]);

	private static readonly MethodInfo _thingCompTickRareMethod
		= AccessTools.DeclaredMethod(typeof(ThingComp), nameof(ThingComp.CompTickRare));

	private static readonly MethodInfo _thingCompTickLongMethod
		= AccessTools.DeclaredMethod(typeof(ThingComp), nameof(ThingComp.CompTickLong));

	private static readonly MethodInfo _hediffCompPostTickMethod
		= AccessTools.DeclaredMethod(typeof(HediffComp), nameof(HediffComp.CompPostTick), [typeof(float).MakeByRefType()]);

	private static readonly MethodInfo _hediffCompPostTickIntervalMethod = AccessTools.DeclaredMethod(
		typeof(HediffComp),
		nameof(HediffComp.CompPostTickInterval),
		[typeof(float).MakeByRefType(), typeof(int)]);

	private static readonly bool _thingCompTickBaseIsTrivial = IsTrivialVoidMethod(_thingCompTickMethod);
	private static readonly bool _thingCompTickIntervalBaseIsTrivial = IsTrivialVoidMethod(_thingCompTickIntervalMethod);
	private static readonly bool _thingCompTickRareBaseIsTrivial = IsTrivialVoidMethod(_thingCompTickRareMethod);
	private static readonly bool _thingCompTickLongBaseIsTrivial = IsTrivialVoidMethod(_thingCompTickLongMethod);
	private static readonly bool _hediffCompPostTickBaseIsTrivial = IsTrivialVoidMethod(_hediffCompPostTickMethod);
	private static readonly bool _hediffCompPostTickIntervalBaseIsTrivial
		= IsTrivialVoidMethod(_hediffCompPostTickIntervalMethod);

	private struct BasePatchState
	{
		public int CheckedTick;
		public bool HasChecked;
		public bool CanSkip;
	}

	private static BasePatchState _thingCompTickState;
	private static BasePatchState _thingCompTickIntervalState;
	private static BasePatchState _thingCompTickRareState;
	private static BasePatchState _thingCompTickLongState;
	private static BasePatchState _hediffCompPostTickState;
	private static BasePatchState _hediffCompPostTickIntervalState;

	private static readonly FishTable<Type, bool> _thingCompTickTypes = new()
	{
		ValueInitializer = static type => InheritsThingCompTick(type)
	};

	private static readonly FishTable<Type, bool> _thingCompTickIntervalTypes = new()
	{
		ValueInitializer = static type => InheritsThingCompTickInterval(type)
	};

	private static readonly FishTable<Type, bool> _thingCompTickRareTypes = new()
	{
		ValueInitializer = static type => InheritsThingCompTickRare(type)
	};

	private static readonly FishTable<Type, bool> _thingCompTickLongTypes = new()
	{
		ValueInitializer = static type => InheritsThingCompTickLong(type)
	};

	private static readonly FishTable<Type, bool> _hediffCompPostTickTypes = new()
	{
		ValueInitializer = static type => InheritsHediffCompPostTick(type)
	};

	private static readonly FishTable<Type, bool> _hediffCompPostTickIntervalTypes = new()
	{
		ValueInitializer = static type => InheritsHediffCompPostTickInterval(type)
	};

	public sealed class ThingTickPatch : FishPrepatch
	{
		public override bool DefaultState =>
#if V1_6
			true;
#else
			false;
#endif

		public override string? Description { get; }
#if V1_6
			= "Skips the component tick loop when every component inherits the unpatched, empty base CompTick method.";
#else
			= "Skips inherited empty ThingComp.CompTick calls when the base method is still unpatched and trivial.";
#endif

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(ThingWithComps), "Tick");

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(ThingWithComps instance)
		{
#if V1_6
			var comps = instance.comps;
			if (comps is null)
				return;
#else
			var comps = instance.AllComps;
#endif

			var count = comps.Count;
#if V1_6
			// Small lists cost less to dispatch normally than to validate a cache.
			if (count >= 4 && CanSkipInheritedBaseThingCompTick())
			{
				ref var cache = ref instance.ThingTickCache();
				if (!ReferenceEquals(cache.Source, comps) || cache.Version != comps._version)
					UpdateThingTickCache(comps, ref cache);

				if (cache.AllTrivial)
					return;
			}
#else
			var canSkipInheritedBase = CanSkipInheritedBaseThingCompTick();
#endif
			for (var i = 0; i < count; i++)
			{
#if V1_6
				instance.comps![i].CompTick();
#else
				var comp = comps[i];
				if (!canSkipInheritedBase || !_thingCompTickTypes.GetOrAdd(comp.GetType()))
					comp.CompTick();
#endif
			}
		}

#if V1_6
		[MethodImpl(MethodImplOptions.NoInlining)]
		private static void UpdateThingTickCache(List<ThingComp> comps, ref ThingTickCache cache)
		{
			var allTrivial = true;
			for (var i = 0; i < comps.Count; i++)
			{
				var comp = comps[i];
				if (comp is null || !_thingCompTickTypes.GetOrAdd(comp.GetType()))
				{
					allTrivial = false;
					break;
				}
			}

			cache.Source = comps;
			cache.Version = comps._version;
			cache.AllTrivial = allTrivial;
		}
#endif
	}

	public sealed class ThingTickIntervalPatch : FishPrepatch
	{
		public override bool DefaultState => false;

		public override string? Description { get; }
			= "Skips inherited empty ThingComp.CompTickInterval calls when the base method is still unpatched and trivial.";

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(ThingWithComps), "TickInterval", [typeof(int)]);

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(ThingWithComps instance, int delta)
		{
		#if V1_6
			var comps = instance.comps;
			if (comps is null)
				return;
		#else
			var comps = instance.AllComps;
		#endif
			var count = comps.Count;
			var canSkipInheritedBase = CanSkipInheritedBaseThingCompTickInterval();
			for (var i = 0; i < count; i++)
			{
		#if V1_6
				var comp = instance.comps![i];
		#else
				var comp = comps[i];
		#endif
				if (!canSkipInheritedBase || !_thingCompTickIntervalTypes.GetOrAdd(comp.GetType()))
					comp.CompTickInterval(delta);
			}
		}
	}

	public sealed class ThingTickRarePatch : FishPrepatch
	{
		public override bool DefaultState => false;

		public override string? Description { get; }
			= "Skips inherited empty ThingComp.CompTickRare calls when the base method is still unpatched and trivial.";

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(ThingWithComps), nameof(ThingWithComps.TickRare));

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(ThingWithComps instance)
		{
		#if V1_6
			var comps = instance.comps;
			if (comps is null)
				return;
		#else
			var comps = instance.AllComps;
		#endif
			var count = comps.Count;
			var canSkipInheritedBase = CanSkipInheritedBaseThingCompTickRare();
			for (var i = 0; i < count; i++)
			{
		#if V1_6
				var comp = instance.comps![i];
		#else
				var comp = comps[i];
		#endif
				if (!canSkipInheritedBase || !_thingCompTickRareTypes.GetOrAdd(comp.GetType()))
					comp.CompTickRare();
			}
		}
	}

	public sealed class ThingTickLongPatch : FishPrepatch
	{
		public override bool DefaultState => false;

		public override string? Description { get; }
			= "Skips inherited empty ThingComp.CompTickLong calls when the base method is still unpatched and trivial.";

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(ThingWithComps), nameof(ThingWithComps.TickLong));

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(ThingWithComps instance)
		{
		#if V1_6
			var comps = instance.comps;
			if (comps is null)
				return;
		#else
			var comps = instance.AllComps;
		#endif
			var count = comps.Count;
			var canSkipInheritedBase = CanSkipInheritedBaseThingCompTickLong();
			for (var i = 0; i < count; i++)
			{
		#if V1_6
				var comp = instance.comps![i];
		#else
				var comp = comps[i];
		#endif
				if (!canSkipInheritedBase || !_thingCompTickLongTypes.GetOrAdd(comp.GetType()))
					comp.CompTickLong();
			}
		}
	}

	public sealed class HediffPostTickPatch : FishPrepatch
	{
		public override bool DefaultState => false;

		public override string? Description { get; }
			= "Skips inherited empty HediffComp.CompPostTick calls when the base method is still unpatched and trivial.";

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(HediffWithComps), nameof(HediffWithComps.PostTick));

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(HediffWithComps instance)
		{
			var comps = instance.comps;
			if (comps == null)
				return;

			var severityAdjustment = 0f;
			var canSkipInheritedBase = CanSkipInheritedBaseHediffCompPostTick();

			for (var i = 0; i < instance.comps!.Count; i++)
			{
				var comp = instance.comps[i];
				if (!canSkipInheritedBase || !_hediffCompPostTickTypes.GetOrAdd(comp.GetType()))
					comp.CompPostTick(ref severityAdjustment);
			}

			if (severityAdjustment != 0f)
				instance.Severity += severityAdjustment;
		}
	}

	public sealed class HediffPostTickIntervalPatch : FishPrepatch
	{
		public override bool DefaultState => false;

		public override string? Description { get; }
			= "Skips inherited empty HediffComp.CompPostTickInterval calls when the base method is still unpatched and trivial.";

		public override MethodBase TargetMethodBase { get; }
			= AccessTools.DeclaredMethod(typeof(HediffWithComps), nameof(HediffWithComps.PostTickInterval), [typeof(int)]);

		public override void Transpiler(ILProcessor ilProcessor, ModuleDefinition module)
			=> ilProcessor.ReplaceBodyWith(ReplacementBody);

		public static void ReplacementBody(HediffWithComps instance, int delta)
		{
			var comps = instance.comps;
			if (comps == null)
				return;

			var severityAdjustment = 0f;
			var canSkipInheritedBase = CanSkipInheritedBaseHediffCompPostTickInterval();

			for (var i = 0; i < instance.comps!.Count; i++)
			{
				var comp = instance.comps[i];
				if (!canSkipInheritedBase || !_hediffCompPostTickIntervalTypes.GetOrAdd(comp.GetType()))
					comp.CompPostTickInterval(ref severityAdjustment, delta);
			}

			if (severityAdjustment != 0f)
				instance.Severity += severityAdjustment;
		}
	}

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseThingCompTick()
		=> CanSkipBaseMethod(_thingCompTickMethod, _thingCompTickBaseIsTrivial, ref _thingCompTickState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseThingCompTickInterval()
		=> CanSkipBaseMethod(_thingCompTickIntervalMethod, _thingCompTickIntervalBaseIsTrivial,
			ref _thingCompTickIntervalState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseThingCompTickRare()
		=> CanSkipBaseMethod(_thingCompTickRareMethod, _thingCompTickRareBaseIsTrivial, ref _thingCompTickRareState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseThingCompTickLong()
		=> CanSkipBaseMethod(_thingCompTickLongMethod, _thingCompTickLongBaseIsTrivial, ref _thingCompTickLongState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseHediffCompPostTick()
		=> CanSkipBaseMethod(_hediffCompPostTickMethod, _hediffCompPostTickBaseIsTrivial,
			ref _hediffCompPostTickState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipInheritedBaseHediffCompPostTickInterval()
		=> CanSkipBaseMethod(_hediffCompPostTickIntervalMethod, _hediffCompPostTickIntervalBaseIsTrivial,
			ref _hediffCompPostTickIntervalState);

	[MethodImpl(MethodImplOptions.AggressiveInlining)]
	private static bool CanSkipBaseMethod(MethodBase method, bool isTrivial, ref BasePatchState state)
	{
		if (!isTrivial)
			return false;

		var tick = Find.TickManager?.TicksGame ?? int.MinValue;
		if (!state.HasChecked || state.CheckedTick != tick)
		{
			state.CanSkip = !HasHarmonyPatches(method);
			state.CheckedTick = tick;
			state.HasChecked = true;
		}

		return state.CanSkip;
	}

	private static bool InheritsThingCompTick(Type type)
		=> AccessTools.Method(type, nameof(ThingComp.CompTick))?.DeclaringType == typeof(ThingComp);

	private static bool InheritsThingCompTickInterval(Type type)
		=> AccessTools.Method(type, nameof(ThingComp.CompTickInterval), [typeof(int)])?.DeclaringType == typeof(ThingComp);

	private static bool InheritsThingCompTickRare(Type type)
		=> AccessTools.Method(type, nameof(ThingComp.CompTickRare))?.DeclaringType == typeof(ThingComp);

	private static bool InheritsThingCompTickLong(Type type)
		=> AccessTools.Method(type, nameof(ThingComp.CompTickLong))?.DeclaringType == typeof(ThingComp);

	private static bool InheritsHediffCompPostTick(Type type)
		=> AccessTools.Method(type, nameof(HediffComp.CompPostTick), [typeof(float).MakeByRefType()])?.DeclaringType
			== typeof(HediffComp);

	private static bool InheritsHediffCompPostTickInterval(Type type)
		=> AccessTools.Method(type, nameof(HediffComp.CompPostTickInterval),
				[typeof(float).MakeByRefType(), typeof(int)])?.DeclaringType
			== typeof(HediffComp);

	private static bool IsTrivialVoidMethod(MethodInfo method)
	{
		if (method.ReturnType != typeof(void)
			|| method.GetMethodBody() is not { ExceptionHandlingClauses.Count: 0, LocalVariables.Count: 0 } body)
		{
			return false;
		}

		var il = body.GetILAsByteArray();
		if (il == null || il.Length == 0)
			return false;

		var sawRet = false;

		for (var i = 0; i < il.Length; i++)
		{
			switch (il[i])
			{
				case 0x00:
					continue;
				case 0x2A when i == il.Length - 1 && !sawRet:
					sawRet = true;
					continue;
				default:
					return false;
			}
		}

		return sawRet;
	}

	private static bool HasHarmonyPatches(MethodBase method)
	{
		var patches = Harmony.GetPatchInfo(method);
		return patches is not null
			&& (patches.Prefixes.Count != 0
				|| patches.Postfixes.Count != 0
				|| patches.Transpilers.Count != 0
				|| patches.Finalizers.Count != 0);
	}
}