package de.tutorials.training.fx.proxy;

import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.*;
import java.util.Hashtable;
import java.util.function.Consumer;

import static java.util.Arrays.asList;

/**
 * @author Thomas Darimont
 */
public class URLStreamFactoryCustomizer {

    public static void useDedicatedProxyForWebkit(Proxy proxy, String protocols) {

        forceInitializationOfOriginalUrlStreamHandlers();
        tryReplaceOriginalUrlStreamHandlersWithScopeProxyAwareVariants(proxy, protocols);
    }

    private static void tryReplaceOriginalUrlStreamHandlersWithScopeProxyAwareVariants(Proxy proxy, String protocols) {

        try {

            Hashtable handlers = tryExtractInternalHandlerTableFromUrl();
            //System.out.println(handlers);

            Consumer<String> wrapStreamHandlerWithScopedProxyHandler = protocol ->
            {
                URLStreamHandler originalHandler = (URLStreamHandler) handlers.get(protocol);
                handlers.put(protocol, new DelegatingScopedProxyAwareUrlStreamHandler(originalHandler, proxy));
            };

            asList(protocols.split(",")).stream().map(String::trim).filter(s -> !s.isEmpty()).forEach(wrapStreamHandlerWithScopedProxyHandler);
            //System.out.println(handlers);

        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    private static Hashtable tryExtractInternalHandlerTableFromUrl() {

        try {
            Field handlersField = URL.class.getDeclaredField("handlers");
            handlersField.setAccessible(true);
            return (Hashtable) handlersField.get(null);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static void forceInitializationOfOriginalUrlStreamHandlers() {

        try {
            new URL("http://.");
            new URL("https://.");
        } catch (MalformedURLException e) {
            throw new RuntimeException(e);
        }
    }

    static class DelegatingScopedProxyAwareUrlStreamHandler extends URLStreamHandler {

        private static final Method openConnectionMethod;
        private static final Method openConnectionWithProxyMethod;

        static {

            try {
                openConnectionMethod = URLStreamHandler.class.getDeclaredMethod("openConnection", URL.class);
                openConnectionWithProxyMethod = URLStreamHandler.class.getDeclaredMethod("openConnection", URL.class, Proxy.class);

                openConnectionMethod.setAccessible(true);
                openConnectionWithProxyMethod.setAccessible(true);

            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        private final URLStreamHandler delegatee;
        private final Proxy proxy;

        public DelegatingScopedProxyAwareUrlStreamHandler(URLStreamHandler delegatee, Proxy proxy) {

            this.delegatee = delegatee;
            this.proxy = proxy;
        }

        @Override
        protected URLConnection openConnection(URL url) throws IOException {

            try {
                if (isWebKitURLLoaderThread(Thread.currentThread())) {

                    //WebKit requested loading the given url, use provided proxy.
                    return (URLConnection) openConnectionWithProxyMethod.invoke(delegatee, url, proxy);
                }

                //Invoke the standard url handler.
                return (URLConnection) openConnectionMethod.invoke(delegatee, url);

            } catch (IllegalAccessException | InvocationTargetException e) {
                throw new RuntimeException(e);
            }
        }

        private boolean isWebKitURLLoaderThread(Thread thread) {

            StackTraceElement[] st = thread.getStackTrace();

            //TODO Add more robust stack-trace inspection.
            return st.length > 4 && st[4].getClassName().startsWith("com.sun.webkit.network");
        }
    }
}