package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;

import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;

import java.io.IOException;
import java.lang.reflect.Modifier;
import java.util.HashSet;
import java.util.Set;
import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;

/**
 * TTL {@link JavassistTransformlet} for {@link java.util.concurrent.Executor}.
 *
 * @author Jerry Lee (oldratlee at gmail dot com)
 * @author wuwen5 (wuwen.55 at aliyun dot com)
 * @see java.util.concurrent.Executor
 * @see java.util.concurrent.ExecutorService
 * @see java.util.concurrent.ThreadPoolExecutor
 * @see java.util.concurrent.ScheduledThreadPoolExecutor
 * @see java.util.concurrent.Executors
 * @since 2.5.1
 */
public class TtlExecutorTransformlet implements JavassistTransformlet {
    private static final Logger logger = Logger.getLogger(TtlExecutorTransformlet.class);

    private static final String TTL_RUNNABLE_CLASS_NAME = "com.alibaba.ttl.TtlRunnable";
    private static final String TTL_CALLABLE_CLASS_NAME = "com.alibaba.ttl.TtlCallable";

    private static final String RUNNABLE_CLASS_NAME = "java.lang.Runnable";
    private static final String CALLABLE_CLASS_NAME = "java.util.concurrent.Callable";

    private static Set<String> EXECUTOR_CLASS_NAMES = new HashSet<String>();

    static {
        EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ThreadPoolExecutor");
        EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ScheduledThreadPoolExecutor");
    }

    @Override
    public boolean needTransform(String className) {
        return EXECUTOR_CLASS_NAMES.contains(className);
    }

    @Override
    public void doTransform(CtClass clazz) throws NotFoundException, CannotCompileException, IOException {
        for (CtMethod method : clazz.getDeclaredMethods()) {
            updateMethodOfExecutorClass(clazz, method);
        }
    }

    private void updateMethodOfExecutorClass(final CtClass clazz, final CtMethod method) throws NotFoundException, CannotCompileException {
        if (method.getDeclaringClass() != clazz) {
            return;
        }
        final int modifiers = method.getModifiers();
        if (!Modifier.isPublic(modifiers) || Modifier.isStatic(modifiers)) {
            return;
        }

        CtClass[] parameterTypes = method.getParameterTypes();
        StringBuilder insertCode = new StringBuilder();
        for (int i = 0; i < parameterTypes.length; i++) {
            CtClass paraType = parameterTypes[i];
            if (RUNNABLE_CLASS_NAME.equals(paraType.getName())) {
                String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TTL_RUNNABLE_CLASS_NAME, i + 1);
                logger.info("insert code before method " + Utils.signatureOfMethod(method) + " of class " + method.getDeclaringClass().getName() + ": " + code);
                insertCode.append(code);
            } else if (CALLABLE_CLASS_NAME.equals(paraType.getName())) {
                String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TTL_CALLABLE_CLASS_NAME, i + 1);
                logger.info("insert code before method " + Utils.signatureOfMethod(method) + " of class " + method.getDeclaringClass().getName() + ": " + code);
                insertCode.append(code);
            }
        }
        if (insertCode.length() > 0) {
            method.insertBefore(insertCode.toString());
        }
    }

}

@@ -63,16 +63,81 @@
 .impl;%0A%0A
+import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;%0A
 import c
@@ -428,19 +428,19 @@
 til.Hash
-Set
+Map
 ;%0Aimport
@@ -450,16 +450,20 @@
 va.util.
+Hash
 Set;%0Aimp
@@ -470,64 +470,43 @@
 ort 
-com.alibaba.ttl.threadpool.agent.internal.logging.Logger
+java.util.Map;%0Aimport java.util.Set
 ;%0A%0A/
@@ -1121,33 +1121,28 @@
 tic 
-final String TTL_RUNNABLE
+Set%3CString%3E EXECUTOR
 _CLA
@@ -1152,131 +1152,34 @@
 NAME
+S
  = 
-%22com.alibaba.ttl.TtlRunnable%22;%0A    private static final String TTL_CALLABLE_CLASS_NAME = %22com.alibaba.ttl.TtlCallable%22;%0A
+new HashSet%3CString%3E();
 %0A   
@@ -1204,397 +1204,501 @@
 nal 
+Map%3C
 String
- RUNNABLE_CLASS_NAME = %22java.lang.Runnable%22;%0A    private static final String CALLABLE_CLASS_NAME = %22java.util.concurrent.Callable%22;%0A%0A    private static Set%3CString%3E EXECUTOR_CLASS_NAMES = new HashSet%3CString%3E();%0A%0A    static %7B%0A        EXECUTOR_CLASS_NAMES.add(%22java.util.concurrent.ThreadPoolExecutor%22);%0A        EXECUTOR_CLASS_NAMES.add(%22java.util.concurrent.ScheduledThreadPoolExecutor
+, String%3E PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS = new HashMap%3CString, String%3E();%0A%0A    static %7B%0A        EXECUTOR_CLASS_NAMES.add(%22java.util.concurrent.ThreadPoolExecutor%22);%0A        EXECUTOR_CLASS_NAMES.add(%22java.util.concurrent.ScheduledThreadPoolExecutor%22);%0A%0A        PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put(%22java.lang.Runnable%22, %22com.alibaba.ttl.TtlRunnable%22);%0A        PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put(%22java.util.concurrent.Callable%22, %22com.alibaba.ttl.TtlCallable
 %22);%0A
@@ -2665,24 +2665,34 @@
     
-CtClass
+final String
  para
+m
 Type
+Name
  = p
@@ -2707,16 +2707,26 @@
 Types%5Bi%5D
+.getName()
 ;%0A      
@@ -2739,53 +2739,74 @@
 if (
-RUNNABLE_CLASS_NAME.equals
+PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.containsKey
 (para
+m
 Type
-.get
 Name
-()
 )) %7B
@@ -2896,425 +2896,67 @@
  1, 
-TTL_RUNNABLE_CLASS_NAME, i + 1);%0A                logger.info(%22insert code before method %22 + Utils.signatureOfMethod(method) + %22 of class %22 + method.getDeclaringClass().getName() + %22: %22 + code);%0A                insertCode.append(code);%0A            %7D else if (CALLABLE_CLASS_NAME.equals
+PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.get
 (para
+m
 Type
-.getName())) %7B%0A                String code = String.format(%22$%25d = %25s.get($%25d, false, true);%22, i + 1, TTL_CALLABLE_CLASS_NAME
+Name)
 , i 
@@ -3298,11 +3298,10 @@
 %7D%0A    %7D%0A
-%0A
 %7D%0A

use data table remove identical code from `if`s #104