using dnlib.DotNet;
using dnlib.DotNet.Emit;
using System;
using System.Collections.Generic;

namespace Unchanger.Obfuscation
{
    public static class Integrity
    {
        public static void Protect(AssemblyDef assembly)
        {
            foreach (ModuleDef module in assembly.Modules)
            {
                Protect(module);
            }
        }

        public static void Protect(ModuleDef module)
        {
            foreach (TypeDef type in module.Types)
            {
                Protect(type);
            }
        }

        public static void Protect(TypeDef type)
        {
            if (type.IsRuntimeSpecialName || type.IsWindowsRuntime)
                return;

            Importer importer = new Importer(type.Module);
            ITypeDefOrRef intRef = importer.Import(typeof(Int32));
            FieldDef fieldInjectedArray = new FieldDefUser("bla", new FieldSig(intRef.ToTypeSig()), FieldAttributes.Static | FieldAttributes.Public);
            type.Fields.Add(fieldInjectedArray);

            int total = 0;
            foreach (MethodDef method in type.Methods)
            {
                Protect(method, fieldInjectedArray,ref  total);
            }
        }

        public static int Protect(MethodDef method, FieldDef field, ref int total)
        {
            if (method == null)
                return 0;

            if (!method.HasBody)
                return 0;

            if (!method.Body.HasInstructions)
                return 0;

            // todo: we can't possibly know the order of the methods executed, so I guess we'll have to
            // use a separate integer for each method.....

            IList<Instruction> instructions = method.Body.Instructions;
            for (int i = 0; i < instructions.Count; i++)
            {
                if (instructions[i].OpCode == OpCodes.Call)
                {
                    MethodDef target = instructions[i].Operand as MethodDef;
                    if (target != method)
                    {
                        // todo: Check if any changes were made, because the number may really be 0
                        int newValue = Protect(target, field, ref total);
                        if (newValue != 0)
                            total = newValue;
                    }
                }
                else if (instructions[i].OpCode == OpCodes.Ldc_I4)
                {
                    int value = (int)instructions[i].Operand;
                    int difference = total - value;

                    instructions[i] = new Instruction(OpCodes.Ldsfld, field);

                    // make sure we cant go out of bounds, "clamp" at 0
                    int newIndex = 0;
                    if (i == 0)
                    {
                        newIndex = 0;
                    }
                    else
                    {
                        newIndex = i - 1;
                    }

                    instructions.Insert(newIndex, new Instruction(OpCodes.Stsfld, field));
                    if (difference < 0)
                    {
                        instructions.Insert(newIndex, new Instruction(OpCodes.Add));
                    }
                    else
                    {
                        instructions.Insert(newIndex, new Instruction(OpCodes.Sub));
                    }
                    instructions.Insert(newIndex, new Instruction(OpCodes.Ldc_I4, Math.Abs(difference)));
                    instructions.Insert(newIndex, new Instruction(OpCodes.Ldsfld, field));

                    total = value;

                    i += 4; // collection size changed
                }
            }

            return total;
        }
    }
}