/*
 * Copyright 2017 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.keycloak.testsuite.adapter.servlet.cluster;

import org.keycloak.admin.client.resource.RealmResource;
import org.keycloak.representations.idm.*;
import org.keycloak.testsuite.Retry;
import org.keycloak.testsuite.adapter.page.EmployeeServletDistributable;
import org.keycloak.testsuite.adapter.page.SAMLServlet;
import org.keycloak.testsuite.auth.page.AuthRealm;
import org.keycloak.testsuite.auth.page.login.*;
import org.keycloak.testsuite.page.AbstractPage;
import org.keycloak.testsuite.util.WaitUtils;

import io.undertow.Undertow;
import io.undertow.server.handlers.ResponseCodeHandler;
import io.undertow.server.handlers.proxy.LoadBalancingProxyClient;
import io.undertow.server.handlers.proxy.ProxyHandler;
import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.jboss.arquillian.container.test.api.*;
import org.jboss.arquillian.graphene.page.Page;
import org.jboss.arquillian.test.api.ArquillianResource;
import org.junit.*;

import org.keycloak.testsuite.adapter.AbstractServletsAdapterTest;
import org.keycloak.testsuite.admin.ApiUtil;
import org.keycloak.testsuite.arquillian.annotation.AppServerContainer;

import org.keycloak.testsuite.util.Matchers;
import org.keycloak.testsuite.util.SamlClient;
import org.keycloak.testsuite.util.SamlClient.Binding;
import org.keycloak.testsuite.util.SamlClientBuilder;
import java.net.MalformedURLException;
import java.util.function.BiConsumer;
import org.apache.http.client.methods.HttpGet;
import org.openqa.selenium.TimeoutException;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.support.ui.WebDriverWait;

import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
import static org.keycloak.testsuite.admin.Users.setPasswordFor;
import static org.keycloak.testsuite.arquillian.AppServerTestEnricher.getNearestSuperclassWithAnnotation;
import static org.keycloak.testsuite.auth.page.AuthRealm.DEMO;
import static org.keycloak.testsuite.util.IOUtil.loadRealm;
import static org.keycloak.testsuite.util.URLAssert.assertCurrentUrlStartsWith;

/**
 *
 * @author hmlnarik
 */
public abstract class AbstractSAMLAdapterClusterTest extends AbstractServletsAdapterTest {

    protected static final String NODE_1_NAME = "ha-node-1";
    protected static final String NODE_2_NAME = "ha-node-2";

    protected final String NODE_1_SERVER_NAME = getAppServerId() + "-" + NODE_1_NAME;
    protected final String NODE_2_SERVER_NAME = getAppServerId() + "-" + NODE_2_NAME;

    protected static final int PORT_OFFSET_NODE_REVPROXY = NumberUtils.toInt(System.getProperty("app.server.reverse-proxy.port.offset"), -1);
    protected static final int HTTP_PORT_NODE_REVPROXY = 8080 + PORT_OFFSET_NODE_REVPROXY;
    protected static final int PORT_OFFSET_NODE_1 = NumberUtils.toInt(System.getProperty("app.server.1.port.offset"), -1);
    protected static final int HTTP_PORT_NODE_1 = 8080 + PORT_OFFSET_NODE_1;
    protected static final int PORT_OFFSET_NODE_2 = NumberUtils.toInt(System.getProperty("app.server.2.port.offset"), -1);
    protected static final int HTTP_PORT_NODE_2 = 8080 + PORT_OFFSET_NODE_2;
    protected static final URI NODE_1_URI = URI.create("http://localhost:" + HTTP_PORT_NODE_1);
    protected static final URI NODE_2_URI = URI.create("http://localhost:" + HTTP_PORT_NODE_2);

    @BeforeClass
    public static void checkPropertiesSet() {
        Assume.assumeThat(PORT_OFFSET_NODE_1, not(is(-1)));
        Assume.assumeThat(PORT_OFFSET_NODE_2, not(is(-1)));
        Assume.assumeThat(PORT_OFFSET_NODE_REVPROXY, not(is(-1)));
    }

    protected static void prepareServerDirectory(String targetSubdirectory) throws IOException {
        Path path = Paths.get(System.getProperty("app.server.home"), targetSubdirectory);
        File targetSubdirFile = path.toFile();
        FileUtils.deleteDirectory(targetSubdirFile);
        FileUtils.forceMkdir(targetSubdirFile);
        FileUtils.copyDirectoryToDirectory(Paths.get(System.getProperty("app.server.home"), "standalone", "deployments").toFile(), targetSubdirFile);
    }
    protected LoadBalancingProxyClient loadBalancerToNodes;
    protected Undertow reverseProxyToNodes;

    @ArquillianResource
    protected ContainerController controller;

    @ArquillianResource
    protected Deployer deployer;

    @Page
    LoginActions loginActionsPage;

    @Override
    public void addTestRealms(List<RealmRepresentation> testRealms) {
        testRealms.add(loadRealm("/adapter-test/keycloak-saml/testsaml-behind-lb.json"));
    }

    @Before
    public void prepareReverseProxy() throws Exception {
        loadBalancerToNodes = new LoadBalancingProxyClient().addHost(NODE_1_URI, NODE_1_NAME).setConnectionsPerThread(10);
        reverseProxyToNodes = Undertow.builder().addHttpListener(HTTP_PORT_NODE_REVPROXY, "localhost").setIoThreads(2).setHandler(new ProxyHandler(loadBalancerToNodes, 5000, ResponseCodeHandler.HANDLE_404)).build();
        reverseProxyToNodes.start();
    }

    @After
    public void stopReverseProxy() {
        reverseProxyToNodes.stop();
    }

    @Before
    public void startServer() throws Exception {
        prepareServerDirectory("standalone-" + NODE_1_NAME);
        controller.start(NODE_1_SERVER_NAME);
        prepareWorkerNode(0, Integer.valueOf(System.getProperty("app.server.1.management.port")));
        prepareServerDirectory("standalone-" + NODE_2_NAME);
        controller.start(NODE_2_SERVER_NAME);
        prepareWorkerNode(1, Integer.valueOf(System.getProperty("app.server.2.management.port")));
        deployer.deploy(EmployeeServletDistributable.DEPLOYMENT_NAME);
        deployer.deploy(EmployeeServletDistributable.DEPLOYMENT_NAME + "_2");
    }

    /**
     * Prepares a worker node
     * @param nodeIndex Node index, counting from 0
     * @param managementPort Port for management operations on this node
     * @throws Exception
     */
    protected abstract void prepareWorkerNode(int nodeIndex, Integer managementPort) throws Exception;

    @After
    public void stopServer() {
        controller.stop(NODE_1_SERVER_NAME);
        controller.stop(NODE_2_SERVER_NAME);
    }

    @Override
    public void setDefaultPageUriParameters() {
        super.setDefaultPageUriParameters();

        testRealmSAMLPostLoginPage.setAuthRealm(DEMO);
        loginPage.setAuthRealm(DEMO);
        loginActionsPage.setAuthRealm(DEMO);
    }

    protected void testLogoutViaSessionIndex(URL employeeUrl, boolean forceRefreshAtOtherNode, BiConsumer<SamlClientBuilder, String> logoutFunction) {
        setPasswordFor(bburkeUser, CredentialRepresentation.PASSWORD);

        final String employeeUrlString;
        try {
            URL employeeUrlAtRevProxy = new URL(employeeUrl.getProtocol(), employeeUrl.getHost(), HTTP_PORT_NODE_REVPROXY, employeeUrl.getFile());
            employeeUrlString = employeeUrlAtRevProxy.toString();
        } catch (MalformedURLException ex) {
            throw new RuntimeException(ex);
        }

        SamlClientBuilder builder = new SamlClientBuilder()
          // Go to employee URL at reverse proxy which is set to forward to first node
          .navigateTo(employeeUrlString)

          // process redirection to login page
          .processSamlResponse(Binding.POST).build()
          .login().user(bburkeUser).build()
          .processSamlResponse(Binding.POST).build()

          // Returned to the page
          .assertResponse(Matchers.bodyHC(containsString("principal=bburke")))

          // Update the proxy to forward to the second node.
          .addStep(() -> updateProxy(NODE_2_NAME, NODE_2_URI, NODE_1_URI));

        if (forceRefreshAtOtherNode) {
            // Go to employee URL at reverse proxy which is set to forward to _second_ node now
            builder
              .navigateTo(employeeUrlString)
              .doNotFollowRedirects()
              .assertResponse(Matchers.bodyHC(containsString("principal=bburke")));
        }

        // Logout at the _second_ node
        logoutFunction.accept(builder, employeeUrlString);

        SamlClient samlClient = builder.execute();
        delayedCheckLoggedOut(samlClient, employeeUrlString);

        // Update the proxy to forward to the first node.
        updateProxy(NODE_1_NAME, NODE_1_URI, NODE_2_URI);
        delayedCheckLoggedOut(samlClient, employeeUrlString);
    }

    private void delayedCheckLoggedOut(SamlClient samlClient, String url) {
        Retry.execute(() -> {
          samlClient.execute(
            (client, currentURI, currentResponse, context) -> new HttpGet(url),
            (client, currentURI, currentResponse, context) -> {
              assertThat(currentResponse, Matchers.bodyHC(not(containsString("principal=bburke"))));
              return null;
            }
          );
        }, 10, 300);
    }

    private void logoutViaAdminConsole() {
        RealmResource demoRealm = adminClient.realm(DEMO);
        String bburkeId = ApiUtil.findUserByUsername(demoRealm, "bburke").getId();
        demoRealm.users().get(bburkeId).logout();
        log.infov("Logged out via admin console");
    }

    @Test
    public void testAdminInitiatedBackchannelLogout(@ArquillianResource
      @OperateOnDeployment(value = EmployeeServletDistributable.DEPLOYMENT_NAME) URL employeeUrl) throws Exception {
        testLogoutViaSessionIndex(employeeUrl, false, (builder, url) -> builder.addStep(this::logoutViaAdminConsole));
    }

    @Test
    public void testAdminInitiatedBackchannelLogoutWithAssertionOfLoggedIn(@ArquillianResource
      @OperateOnDeployment(value = EmployeeServletDistributable.DEPLOYMENT_NAME) URL employeeUrl) throws Exception {
        testLogoutViaSessionIndex(employeeUrl, true, (builder, url) -> builder.addStep(this::logoutViaAdminConsole));
    }

    @Test
    public void testUserInitiatedFrontchannelLogout(@ArquillianResource
      @OperateOnDeployment(value = EmployeeServletDistributable.DEPLOYMENT_NAME) URL employeeUrl) throws Exception {
        testLogoutViaSessionIndex(employeeUrl, false, (builder, url) -> {
            builder
              .navigateTo(url + "?GLO=true")
              .processSamlResponse(Binding.POST).build()    // logout request
              .processSamlResponse(Binding.POST).build()    // logout response
            ;
        });
    }

    @Test
    public void testUserInitiatedFrontchannelLogoutWithAssertionOfLoggedIn(@ArquillianResource
      @OperateOnDeployment(value = EmployeeServletDistributable.DEPLOYMENT_NAME) URL employeeUrl) throws Exception {
        testLogoutViaSessionIndex(employeeUrl, true, (builder, url) -> {
            builder
              .navigateTo(url + "?GLO=true")
              .processSamlResponse(Binding.POST).build()    // logout request
              .processSamlResponse(Binding.POST).build()    // logout response
            ;
        });
    }

    protected void updateProxy(String hostToPointToName, URI hostToPointToUri, URI hostToRemove) {
        loadBalancerToNodes.removeHost(hostToRemove);
        loadBalancerToNodes.addHost(hostToPointToUri, hostToPointToName);
        log.infov("Reverse proxy will direct requests to {0}", hostToPointToUri);
    }

    protected void assertSuccessfulLogin(SAMLServlet page, UserRepresentation user, Login loginPage, String expectedString) {
        page.navigateTo();
        assertCurrentUrlStartsWith(loginPage);
        loginPage.form().login(user);
        WebDriverWait wait = new WebDriverWait(driver, WaitUtils.PAGELOAD_TIMEOUT_MILLIS / 1000);
        wait.until((WebDriver d) -> d.getPageSource().contains(expectedString));
    }

    protected void delayedCheckLoggedOut(AbstractPage page, AuthRealm loginPage) {
        Retry.execute(() -> {
            try {
                checkLoggedOut(page, loginPage);
            } catch (AssertionError | TimeoutException ex) {
                driver.navigate().refresh();
                log.debug("[Retriable] Timed out waiting for login page");
                throw new RuntimeException(ex);
            }
        }, 10, 100);
    }

    protected void checkLoggedOut(AbstractPage page, AuthRealm loginPage) {
        page.navigateTo();
        WaitUtils.waitForPageToLoad();
        assertCurrentUrlStartsWith(loginPage);
    }

    private String getAppServerId() {
        Class<?> annotatedClass = getNearestSuperclassWithAnnotation(this.getClass(), AppServerContainer.class);

        return (annotatedClass == null ? "<cannot-find-@AppServerContainer>"
                : annotatedClass.getAnnotation(AppServerContainer.class).value());
    }

}
