﻿using System.Linq;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Storm.Inject.Params;
using Storm.Util;

namespace Storm.Inject.Cecil
{
    public class CecilMethodCallDetourInjector : Injector
    {
        private readonly AssemblyDefinition def;
        private readonly MethodCallDetourParams @params;
        private readonly AssemblyDefinition self;

        public CecilMethodCallDetourInjector(AssemblyDefinition self, AssemblyDefinition def, MethodCallDetourParams @params)
        {
            this.self = self;
            this.def = def;
            this.@params = @params;
        }

        public void Init()
        {
        }

        public void Inject()
        {
            var from = def.GetTypeRef(@params.FromClass) ?? self.GetTypeRef(@params.FromClass, true);
            if (from == null)
            {
                Logging.FDbgLog(LogLevel.Warn, "[{0}] Unable to find from type", GetType().Name);
                Logging.FDbgLog(LogLevel.Warn, "\t{0} {1}", @params.FromClass, @params.ToClass);
                return;
            }

            var to = def.GetTypeRef(@params.ToClass) ?? self.GetTypeRef(@params.ToClass, true);
            if (to == null)
            {
                Logging.FDbgLog(LogLevel.Warn, "[{0}] Unable to find to type", GetType().Name);
                Logging.FDbgLog(LogLevel.Warn, "\t{0} {1}", @params.FromClass, @params.ToClass);
                return;
            }

            var methods = def.Modules
                .SelectMany(m => m.Types)
                .Where(t => @params.IgnoredTypes == null || !@params.IgnoredTypes.Contains(t.FullName))
                .SelectMany(t => t.Methods)
                .Where(m => m.HasBody).ToList();

            foreach (var method in methods)
            {
                var body = method.Body;
                var instructions = body.Instructions;
                var processor = body.GetILProcessor();

                for (int i = 0; i < instructions.Count; i++)
                {
                    Instruction ins = instructions[i];
                    if (ins.OpCode == OpCodes.Call 
                        || ins.OpCode == OpCodes.Callvirt 
                        || ins.OpCode == OpCodes.Calli)
                    {
                        var @ref = (MethodReference)ins.Operand;
                        if (@ref.DeclaringType.FullName.Equals(@params.FromClass))
                        {
                            var requiredParamCount = @ref.Parameters.Count;
                            if (!@ref.Resolve().IsStatic)
                            {
                                requiredParamCount += 1;
                            }

                            var redirect = to
                                .Resolve()
                                .Methods
                                .FirstOrDefault(m => m.Name.Equals(@ref.Name) && requiredParamCount == m.Parameters.Count);

                            if (redirect != null)
                            {
                                var call = def.Import(redirect);
                                var gim = @ref as GenericInstanceMethod;
                                if (gim != null)
                                {
                                    var genericCall = new GenericInstanceMethod(def.Import(redirect));
                                    foreach (var arg in gim.GenericArguments)
                                    {
                                        genericCall.GenericArguments.Add(arg);
                                    }
                                    call = genericCall;
                                }
                                
                                processor.Replace(ins, processor.Create(OpCodes.Call, call));
                            }
                        }
                    }
                }
            }
        }

        public object GetParams()
        {
            return @params;
        }
    }
}