package quant.rich.emoney.patch.okhttp;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.asm.Advice;
import okhttp3.Request;
import quant.rich.emoney.util.SpringContextHolder;
import net.bytebuddy.dynamic.loading.ClassReloadingStrategy;
import static net.bytebuddy.matcher.ElementMatchers.*;

import java.io.Serializable;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.function.Consumer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.SpringApplication;
import org.springframework.context.ConfigurableApplicationContext;

public class PatchOkHttp {
    private static final Random random = new Random();
    private static final Logger log = LoggerFactory.getLogger(PatchOkHttp.class);
    private static final Set<PatchOkHttpRule> rules = new CopyOnWriteArraySet<>();
    
    private static boolean isHooked = false;

    /**
     * 应用指定 Patch 规则
     * @param rule
     * @return 如果 rule 未设置 id，则生成随机 id 并返回，否则返回 rule.getId()
     */
    public static Integer apply(PatchOkHttpRule rule) {
        
        if (rule.getId() == null) {
            Integer[] randomIds = new Integer[] { random.nextInt() };
            while (rules.parallelStream().anyMatch(r -> Integer.compare(randomIds[0], r.getId()) == 0)) {
                // 如果 rules 有已有 ID 的话要重新设置 ID，不能重复
                randomIds[0] = random.nextInt();
            }
            rule.setId(randomIds[0]);
            log.debug("PatchOkHttp.apply(rule.id={})", randomIds[0]);
        }
        
        rules.add(rule);
        //log.debug("PatchOkHttp.apply() 在 ClassLoader {} 中运行", PatchOkHttp.class.getClassLoader());
        if (!isHooked) hook();
        return rule.getId();
    }
    
    public static Serializable apply(PatchOkHttpRule.Builder builder) {
        return PatchOkHttp.apply(builder.build());
    }
    
    public static Serializable apply(Consumer<PatchOkHttpRule.Builder> r) {
        PatchOkHttpRule.Builder builder = PatchOkHttpRule.when();
        r.accept(builder);
        return PatchOkHttp.apply(builder.build());
    }
    
    public static void match(RequestContext ctx, String currentHeader, Consumer<String> consumer) {
        // log.debug("PatchOkHttp.match() 在 ClassLoader {} 中运行", PatchOkHttp.class.getClassLoader());
        for (PatchOkHttpRule rule : PatchOkHttp.rules) {
            if (rule.matches(ctx)) {
                log.debug("PatchOkHttp.match() 匹配到规则 rule.id={}", rule.getId());
                rule.apply(ctx, currentHeader, consumer);
            }
        }
    }

    /**
     * <p>通过指定 Advice 重定义指定类的指定方法<br>
     * 此处为：<code>okhttp3.Request$Builder.header(String, String)</code></p>
     * <p>
     * <b>注意：</b>必须保证欲重定义的类的方法的 ClassLoader 和调用该方法的 ClassLoader
     * 一致，否则会导致两者的静态类的静态变量不一致，<b>尤其是项目引入了 spring-boot-devtool
     * 时</b>，一般调用该方法的类由 RestartClassLoader 加载，而 okhttp 包由于是 jar 包，一般交由默认
     * ClassLoaders 进行加载，这就导致 PatchOkHttp 在两个 ClassLoader 中不一致，实际注入方法内 <code>PatchOkHttp.<i><b>rules</b></i></code>
     * 为空，无论如何都无法命中设定的规则。</p>
     * <p>解决方法：<ol>
     * <li>在 src/main/resources 下建立 META-INF 目录，新建
     * spring-devtools.properties，并在其内增加一行“restart.include.okhttp3=.*okhttp-.*\.jar”，指定由
     * RestartClassLoader 加载 okhttp-*.jar 包；</li>
     * <li>在 application.yml 中，增加需要额外扫描的路径，即设置参数“spring.devtools.restart.additional-pathes:
     * lib/”，否则即便是进行了第一步，也不会生效；</li>
     * <li>在方法体内，使用外部 RestartClassLoader 载入欲重定义的方法类，并在调用 ByteBuddy 时指定该方法类和 ClassLoader。</li>
     * </ol>
     * </p>
     * <p>进一步地，如果注入很重要，请判断欲重定义类所属 ClassLoader 和外部 ClassLoader 是否一致，如果不一致的，终止程序。</p>
     * @see org.springframework.boot.devtools.restart.classloader.RestartClassLoader
     * @see jdk.internal.loader.ClassLoaders
     */
    private static void hook() {
        try {

            ByteBuddyAgent.install();
            
            ClassLoader appLoader = PatchOkHttp.class.getClassLoader();
            Class<?> builderClass = appLoader.loadClass("okhttp3.Request$Builder");
            

            if (builderClass.getClassLoader() != appLoader) {
                int exitCode = SpringApplication.exit(
                    SpringContextHolder.getBean(ConfigurableApplicationContext.class), () -> 0);
                System.exit(exitCode);
            }

            new ByteBuddy().redefine(builderClass)
                    .visit(Advice.to(HeaderInterceptor.class)
                            .on(named("header").and(takesArguments(String.class, String.class))))
                    .make().load(appLoader, ClassReloadingStrategy.fromInstalledAgent());
            
            
            isHooked = true;
            log.debug("[PatchOkHttp] Delegation hook installed.");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static class HeaderInterceptor {


        public static final Logger log = LoggerFactory.getLogger(HeaderInterceptor.class);
        @Advice.OnMethodEnter(
                skipOn = Advice.OnNonDefaultValue.class)
        public static boolean intercept(
                @Advice.This Request.Builder builder,
                @Advice.Argument(0) String name,
                @Advice.Argument(value = 1, readOnly = false) String value,
                @Advice.Local("call") Callable<okhttp3.Request.Builder> superCall
        ) throws Exception {
            
            try {
                // 获取上下文
                RequestContext ctx = new RequestContext(builder);
                
                String[] modified = new String[]{value};
                Consumer<String> setter = new ModifiedHeaderValueSetter(modified);
                PatchOkHttp.match(ctx, name, setter);

                // 替换原参数值
                value = modified[0];

            } catch (Exception e) {
                e.printStackTrace();
            }
            return false;
        }
    }
    
    public static class ModifiedHeaderValueSetter implements Consumer<String> {
        private final String[] holder;
        public ModifiedHeaderValueSetter(String[] holder) {
            this.holder = holder;
        }
        @Override
        public void accept(String newVal) {
            holder[0] = newVal;
        }
    }
}