package org.openslx.dozmod.util;

import java.io.FileInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class CombinedTrustManager {

	private final static Logger LOGGER = LogManager.getLogger(CombinedTrustManager.class);

	private static SSLContext sslContext = null;

	private static CombinedX509TrustManager delegatingTrustManager = null;

	public static void install() {
		// On Windows, use system store in addition to the Java one
		List<X509TrustManager> managers = new ArrayList<>();
		char[] password = "changeit".toCharArray();

		LOGGER.info("Installing Fallback X509 truster");

		try {
			// --- Load Java default trust store (cacerts) ---
			String javaHome = System.getProperty("java.home");
			String cacertsPath = javaHome + "/lib/security/cacerts";
			KeyStore javaTrustStore = KeyStore.getInstance("JKS");

			try (FileInputStream fis = new FileInputStream(cacertsPath)) {
				javaTrustStore.load(fis, password);
			}
			addKeyStore(managers, javaTrustStore, "Java");
		} catch (Exception e) {
			LOGGER.warn("Error adding java certificate store", e);
		}

		if (OsHelper.isWindows()) {
			try {
				// --- Load Windows root store ---
				KeyStore winRootStore = KeyStore.getInstance("Windows-ROOT");
				winRootStore.load(null, null);
				addKeyStore(managers, winRootStore, "Windows-ROOT");
			} catch (Exception e) {
				LOGGER.warn("Error adding Windows-ROOT certificate store", e);
			}
			try {
				// --- Load Windows user store ---
				KeyStore winUserStore = KeyStore.getInstance("Windows-MY");
				winUserStore.load(null, null);
				addKeyStore(managers, winUserStore, "Windows-MY");
			} catch (Exception e) {
				LOGGER.warn("Error adding Windows-MY certificate store", e);
			}
		}

		try {
			KeyStore shippedStore = KeyStore.getInstance("JKS");
			try (InputStream is = ResourceLoader.getStream("/data/truststore.jks")) {
				shippedStore.load(is, password);
			}
			addKeyStore(managers, shippedStore, "Shipped");
		} catch (Exception e) {
			LOGGER.warn("Error adding shipped certificate store", e);
		}

		if (managers.isEmpty()) {
			LOGGER.warn("Couldn't load any trust manager - using default one");
			return;
		}

		try {
			// --- Combine using delegating trust manager ---
			delegatingTrustManager = new CombinedX509TrustManager(managers);

			sslContext = SSLContext.getInstance("TLS");
			sslContext.init(null, getTrustManagers(), null);
			SSLContext.setDefault(sslContext);
			HttpsURLConnection.setDefaultSSLSocketFactory(sslContext.getSocketFactory());
		} catch (Exception e) {
			LOGGER.warn("Error installing custom trust manager", e);
		}
	}

	private static void addKeyStore(List<X509TrustManager> list, KeyStore store, String name)
			throws Exception {
		LOGGER.info(name + " entries: " + store.size());
		if (store.size() == 0)
			return; // Empty ones cause problems
		TrustManagerFactory javaTMF = TrustManagerFactory.getInstance(
				TrustManagerFactory.getDefaultAlgorithm());
		javaTMF.init(store);
		X509TrustManager javaTrustManager = getX509TrustManager(javaTMF);
		list.add(javaTrustManager);
	}

	public static TrustManager getTrustManager() {
		return delegatingTrustManager;
	}

	public static TrustManager[] getTrustManagers() {
		if (delegatingTrustManager == null)
			return null;
		return new TrustManager[] { delegatingTrustManager };
	}

	// Extract the first X509TrustManager from the factory
	private static X509TrustManager getX509TrustManager(TrustManagerFactory tmf) throws Exception {
		for (TrustManager tm : tmf.getTrustManagers()) {
			if (tm instanceof X509TrustManager) {
				return (X509TrustManager) tm;
			}
		}
		throw new IllegalStateException("No X509TrustManager found");
	}

	// Delegating trust manager implementation
	public static class CombinedX509TrustManager implements X509TrustManager {
		private final List<X509TrustManager> managers;
		private X509Certificate[] issuers = null;

		public CombinedX509TrustManager(List<X509TrustManager> managers) {
			this.managers = managers;
		}

		@Override
		public void checkClientTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			CertificateException cached = null;

			for (X509TrustManager tm : managers) {
				try {
					tm.checkClientTrusted(chain, authType);
					return;
				} catch (CertificateException e) {
					cached = e;
				} catch (RuntimeException rte) {
					LOGGER.warn("Other exception in checkClientTrusted", rte);
				}
			}
			if (cached != null)
				throw cached;
			throw new CertificateException("Unknown exception in combined trust manager");
		}

		@Override
		public void checkServerTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			CertificateException cached = null;

			for (X509TrustManager tm : managers) {
				try {
					tm.checkServerTrusted(chain, authType);
					return;
				} catch (CertificateException e) {
					cached = e;
				} catch (RuntimeException rte) {
					LOGGER.warn("Other exception in checkServerTrusted", rte);
				}
			}
			if (cached != null)
				throw cached;
			throw new CertificateException("Unknown exception in combined trust manager");
		}

		@Override
		public X509Certificate[] getAcceptedIssuers() {
			if (issuers == null) {
				Set<X509Certificate> certs = new HashSet<>();
				for (X509TrustManager tm : managers) {
					try {
						certs.addAll(Arrays.asList(tm.getAcceptedIssuers()));
					} catch (Exception e) {
						LOGGER.warn("Error adding accepted issuers to combined return value", e);
					}
				}
				issuers = certs.toArray(new X509Certificate[certs.size()]);
			}
			return issuers;
		}
	}

}