using System;
using System.Linq;
using dnlib.DotNet.Emit;
using dnlib.DotNet.PolymorphicWriter;
using dnlib.DotNet.Writer;
using HarmonyLib;
using Obfuz.Settings;

internal static class DnlibPolymorphicContext
{
    [ThreadStatic] public static byte[] OpcodeMapping;
}

[HarmonyPatch(typeof(PolymorphicMetadata), "WriteMethodBodies")]
static class WriteMethodBodiesPatch
{
    static void Prefix()
    {
        if (!KedrExtensionSettings.Instance.enableOpcodeShuffle)
            return;

        var seed = HashUtil.Hash(ObfuzSettings.Instance.polymorphicDllSettings.codeGenerationSecretKey);
        var opcodeMapping = Enumerable.Range(0, 248).Select(x => (byte)x).ToArray();
        var rng = new System.Random((int)seed);
        RandomUtil.Shuffle(opcodeMapping, rng);

        DnlibPolymorphicContext.OpcodeMapping = opcodeMapping;
    }

    static void Finalizer()
    {
        DnlibPolymorphicContext.OpcodeMapping = null;
    }
}

[HarmonyPatch(typeof(MethodBodyWriterBase), "WriteOpCode")]
static class WriteOpCodePatch
{
    static bool Prefix(ref ArrayWriter writer, Instruction instr)
    {
        if (DnlibPolymorphicContext.OpcodeMapping is null)
            return true;

        WriteOpCode(ref writer, instr);
        return false;
    }

    private static void WriteOpCode(ref ArrayWriter writer, Instruction instr)
    {
        var code = (int)instr.OpCode.Code;
        int num = code >> 8;

        if (code <= 255)
        {
            if (code < (int)dnlib.DotNet.Emit.Code.Prefix7)
            {
                code = DnlibPolymorphicContext.OpcodeMapping[code];
            }
            
            writer.WriteByte((byte)code);
            return;
        }

        switch (num)
        {
            case 240:
            case 241:
            case 242:
            case 243:
            case 244:
            case 245:
            case 246:
            case 247:
            case 248:
            case 249:
            case 250:
            case 251:
            case 254:
                writer.WriteByte((byte)num);
                writer.WriteByte((byte)instr.OpCode.Code);
                return;
        }

        switch (instr.OpCode.Code)
        {
            case dnlib.DotNet.Emit.Code.UNKNOWN1:
                writer.WriteByte(0);
                break;
            case dnlib.DotNet.Emit.Code.UNKNOWN2:
                writer.WriteUInt16(0);
                break;
            default:
                throw new Exception("Unknown opcode");
        }
    }
}