package edu.kit.scc.dei.ecplean;

import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.util.Observable;

import javax.xml.namespace.QName;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathException;
import javax.xml.xpath.XPathExpression;
import javax.xml.xpath.XPathFactory;

import org.apache.hc.client5.http.auth.AuthenticationException;
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
import org.apache.hc.client5.http.classic.methods.HttpPost;
import org.apache.hc.client5.http.impl.auth.BasicScheme;
import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse;
import org.apache.hc.client5.http.impl.classic.HttpClients;
import org.apache.hc.client5.http.protocol.HttpClientContext;
import org.apache.hc.core5.http.HttpHeaders;
import org.apache.hc.core5.http.HttpHost;
import org.apache.hc.core5.http.HttpStatus;
import org.apache.hc.core5.http.ParseException;
import org.apache.hc.core5.http.io.entity.EntityUtils;
import org.apache.hc.core5.http.io.entity.StringEntity;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.w3c.dom.Document;
import org.xml.sax.EntityResolver;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

public abstract class ECPAuthenticatorBase extends Observable {

	protected static Logger logger = LogManager.getLogger(ECPAuthenticatorBase.class);
	
	protected ECPAuthenticationInfo authInfo;
	protected CloseableHttpClient client;
	protected DocumentBuilderFactory documentBuilderFactory;
	protected XPathFactory xpathFactory;
	protected NamespaceResolver namespaceResolver;
	protected TransformerFactory transformerFactory;

	protected boolean retryWithoutAt;

	public ECPAuthenticatorBase(CloseableHttpClient client) {
		this.client = client == null ? HttpClients.createSystem() : client;

		documentBuilderFactory = DocumentBuilderFactory.newInstance();
		documentBuilderFactory.setNamespaceAware(true);
		
		xpathFactory = XPathFactory.newInstance();
		namespaceResolver = new NamespaceResolver();
		namespaceResolver.addNamespace("ecp", "urn:oasis:names:tc:SAML:2.0:profiles:SSO:ecp");
		namespaceResolver.addNamespace("S", "http://schemas.xmlsoap.org/soap/envelope/");
		namespaceResolver.addNamespace("paos", "urn:liberty:paos:2003-08");
		
		transformerFactory = TransformerFactory.newInstance();
	}

	public ECPAuthenticatorBase() {
		this(null);
	}

	private CloseableHttpResponse exec(Document idpRequest, String user, String pass)
			throws ECPAuthenticationException {
		final HttpHost httpHost = HttpHost.create(authInfo.getSpUrl());

		// setup basic authentication
		final UsernamePasswordCredentials userCredentials = new UsernamePasswordCredentials(user, pass.toCharArray());
      final BasicScheme basicAuth = new BasicScheme();
      basicAuth.initPreemptive(userCredentials);

      // create local HTTP context for basic authentication
      final HttpClientContext httpContext = HttpClientContext.create();
      httpContext.resetAuthExchange(httpHost, basicAuth);

      // create POST request to IdP
      final HttpPost httpPost = new HttpPost(authInfo.getIdpEcpEndpoint().toString());

      // fill content of POST request
		try {
			httpPost.setEntity(new StringEntity(documentToString(idpRequest)));
		} catch (TransformerException e1) {
			logger.warn("Error setting XML payload of IdP POST");
			throw new ECPAuthenticationException(e1);
		}

		// set content type of POST request
		httpPost.setHeader(HttpHeaders.CONTENT_TYPE, "text/xml; charset=utf-8");

		// set basic authentication header for POST request
		try {
			httpPost.setHeader(HttpHeaders.AUTHORIZATION, basicAuth.generateAuthResponse(httpHost, httpPost, httpContext));
		} catch (AuthenticationException e) {
			logger.warn("Error setting Authentication header for IdP POST");
			throw new ECPAuthenticationException(e);
		}

		// send POST request to IdP
		try {
			return client.execute(httpPost, httpContext);
		} catch (Exception e) {
			httpPost.reset();
			logger.error("Could not submit PAOS request to IdP");
			throw new ECPAuthenticationException(e);
		}
	}

	protected Document authenticateIdP(Document idpRequest) throws ECPAuthenticationException {
		logger.info("Sending initial IdP Request");

		CloseableHttpResponse httpResponse = null;
		String user = authInfo.getUsername();
		String pass = authInfo.getPassword();
		int at = user.lastIndexOf('@');
		boolean failed = false;
		try {
			httpResponse = exec(idpRequest, user, pass);
			failed = (httpResponse.getCode() == HttpStatus.SC_UNAUTHORIZED);
		} catch (ECPAuthenticationException e) {
			logger.debug("Could not submit PAOS request to IdP");
			if (at == -1)
				throw new ECPAuthenticationException(e);
			failed = true;
		}
		if (at != -1 && failed && retryWithoutAt) {
			// Retrying without the @ in the username is desired
			user = user.substring(0, at);
			try {
				httpResponse = exec(idpRequest, user, pass);
			} catch (ECPAuthenticationException e) {
				logger.debug("Could not submit PAOS request to IdP");
				throw new ECPAuthenticationException(e);
			}
		}
		String responseBody;
		try {
			responseBody = EntityUtils.toString(httpResponse.getEntity());
		} catch (RuntimeException | IOException | ParseException e) {
			logger.debug("Could not read response from IdP");
			throw new ECPAuthenticationException(e);
		}
		try {
			return buildDocumentFromString(responseBody);
		} catch (IOException | SAXException | ParserConfigurationException | RuntimeException e) {
			logger.debug("Could not parse XML response from IdP:\n" + responseBody);
			throw new ECPAuthenticationException(e);
		}

	}
	
	protected Document buildDocumentFromString(String input)
			throws IOException, ParserConfigurationException, SAXException {
		DocumentBuilder builder = documentBuilderFactory.newDocumentBuilder();
		builder.setEntityResolver(new EntityResolver() {
			@Override
			public InputSource resolveEntity(String publicId, String systemId) throws SAXException, IOException {
				return new InputSource(new StringReader(""));
			}
		});
		return builder.parse(new InputSource(new StringReader(input)));
	}

	protected Object queryDocument(Document xmlDocument, String expression,
			QName returnType) throws XPathException {
		XPath xpath = xpathFactory.newXPath();
		xpath.setNamespaceContext(namespaceResolver);
		XPathExpression xPathExpression = xpath.compile(expression);
		return xPathExpression.evaluate(xmlDocument, returnType);
	}

	protected String documentToString(Document xmlDocument)
			throws TransformerConfigurationException, TransformerException {
		Transformer transformer = transformerFactory.newTransformer();

		StreamResult result = new StreamResult(new StringWriter());
		DOMSource source = new DOMSource(xmlDocument);
		transformer.transform(source, result);

		return result.getWriter().toString();
	}

	public CloseableHttpClient getHttpClient() {
		return client;
	}

	public void setRetryWithoutAt(boolean b) {
		this.retryWithoutAt = b;
	}

}