package quant.rich.emoney.patch.okhttp;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.asm.Advice;
import okhttp3.Request;
import net.bytebuddy.dynamic.loading.ClassReloadingStrategy;
import static net.bytebuddy.matcher.ElementMatchers.*;

import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

    public static void apply(PatchOkHttpRule.Builder builder) {
        rules.add(builder.build());
        if (!isHooked) hook();
    }
    
    public static void apply(Consumer<PatchOkHttpRule.Builder> r) {
        PatchOkHttpRule.Builder builder = PatchOkHttpRule.when();
        r.accept(builder);
        rules.add(builder.build());
        if (!isHooked) hook();
    }
    
    public static void match(RequestContext ctx, String currentHeader, Consumer<String> consumer) {
        for (PatchOkHttpRule rule : PatchOkHttp.rules) {
            if (rule.matches(ctx)) {
                rule.apply(ctx, currentHeader, consumer);
            }
        }
    }

    private static void hook() {
        try {
            ByteBuddyAgent.install();

            new ByteBuddy()
              .redefine(Class.forName("okhttp3.Request$Builder"))
              .visit(Advice.to(HeaderInterceptor.class)
                .on(named("header").and(takesArguments(String.class, String.class))))
              .make()
              .load(Class.forName("okhttp3.Request$Builder").getClassLoader(),
                    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 {

            // print调试
            log.debug("[intercept] name = {}, value = {}", name, value);

            try {
                // 获取 headers map
                
                
                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;
        }
    }
}