/*
 * (c) 2016 mike chaberski
 */
package com.github.mike10004.xvfbmanagerexample;

import com.github.mike10004.xvfbmanager.XvfbController;
import com.github.mike10004.xvfbmanager.XvfbManager;
import com.github.mike10004.xvfbselenium.WebDriverSupport;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.Resources;
import com.google.common.net.HttpHeaders;
import com.google.common.net.MediaType;
import io.github.bonigarcia.wdm.WebDriverManager;
import org.apache.commons.lang3.tuple.Pair;
import org.mockserver.integration.ClientAndServer;
import org.mockserver.model.HttpRequest;
import org.mockserver.model.HttpResponse;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.firefox.FirefoxDriver;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.net.URL;

import static com.google.common.base.Preconditions.checkArgument;

public class XvfbManagerExample {

    public static void main(String[] args) throws IOException {
        if (args.length < 0 || !browserMap.containsKey(args[0])) {
            System.err.format("argument must be one of {%s}%n", Joiner.on(", ").join(browserMap.keySet()));
            System.exit(1);
        }
        browseAndCaptureScreenshot(args[0]);
    }

    public static void browseAndCaptureScreenshot(String browserKey) throws IOException {
        ClientAndServer server = ClientAndServer.startClientAndServer(0);
        try {
            String host = "localhost", path = "/";
            int port = server.getPort();
            String html = "<html><head><title>Example</title></head><body><img src=\"cat.jpg\"></body></html>";
            byte[] imageBytes = Resources.toByteArray(XvfbManagerExample.class.getResource("/smiling-cat-in-public-domain.jpg"));
            server.when(HttpRequest.request(path)).respond(HttpResponse.response(html).withHeader(HttpHeaders.CONTENT_TYPE, MediaType.HTML_UTF_8.toString()));
            server.when(HttpRequest.request(path + "cat.jpg")).respond(HttpResponse.response().withBody(imageBytes).withHeader(HttpHeaders.CONTENT_TYPE, MediaType.JPEG.toString()));
            URL url = new URL("http", host, port, path);
            BufferedImage screenshotImage;
            XvfbManager xvfb = new XvfbManager();
            try (XvfbController ctrl = xvfb.start()) {
                screenshotImage = browse(browserKey, url, ctrl.getDisplay());
            }
            System.out.format("captured %dx%d screenshot%n", screenshotImage.getWidth(), screenshotImage.getHeight());
            File screenshotFile = new File(System.getProperty("user.dir")).toPath().resolve("target").resolve("screenshot-" + System.currentTimeMillis() + ".png").toFile();
            ImageIO.write(screenshotImage, "png", screenshotFile);
        } finally {
            server.stop();
        }
    }

    /**
     * Visit a given URL and return a screenshot.
     * @param url the url
     * @return an image captured by the browser
     * @throws IOException
     */
    private static BufferedImage browse(String browserKey, URL url, String display) throws IOException {
        Pair<Class<? extends WebDriver>, Function<String, ? extends WebDriver>> driverStuff = browserMap.get(browserKey);
        checkArgument(driverStuff != null, "unsupported browser: %s", browserKey);
        System.out.format("browsing %s on display %s with %s%n", url, display, driverStuff.getLeft());
        Class<? extends WebDriver> webDriverClass = driverStuff.getLeft();
        WebDriverManager.getInstance(webDriverClass).setup();
        WebDriver webDriver = driverStuff.getRight().apply(display);
        try {
            webDriver.get(url.toString());
            byte[] screenshotBytes = ((TakesScreenshot)webDriver).getScreenshotAs(OutputType.BYTES);
            BufferedImage screenshotImage = ImageIO.read(new ByteArrayInputStream(screenshotBytes));
            return screenshotImage;
        } finally {
            webDriver.quit();
        }
    }

    private static final ImmutableMap<String, Pair<Class<? extends WebDriver>, Function<String, ? extends WebDriver>>> browserMap =
            ImmutableMap.<String, Pair<Class<? extends WebDriver>, Function<String, ? extends WebDriver>>>builder()
            .put("firefox", Pair.of(org.openqa.selenium.firefox.MarionetteDriver.class, new Function<String, FirefoxDriver>(){
                @Override
                public FirefoxDriver apply(String display) {
                    return WebDriverSupport.firefoxOnDisplay(display).create();
                }
            }))
            .put("chrome", Pair.of(ChromeDriver.class, new Function<String, ChromeDriver>(){
                @Override
                public ChromeDriver apply(String display) {
                    return WebDriverSupport.chromeOnDisplay(display).create();
                }
            }))
            .build();
}

@@ -930,22 +930,123 @@
 .Firefox
-Driver
+Binary;%0Aimport org.openqa.selenium.firefox.FirefoxDriver;%0Aimport org.openqa.selenium.firefox.FirefoxProfile
 ;%0A%0Aimpor
@@ -1317,16 +1317,81 @@
 mple %7B%0A%0A
+    public static final String ENV_FIREFOX_BIN = %22FIREFOX_BIN%22;%0A%0A
     publ
@@ -4886,32 +4886,464 @@
 ring display) %7B%0A
+                    String firefoxBin = System.getenv(ENV_FIREFOX_BIN);%0A                    if (firefoxBin != null) %7B%0A                        System.out.format(%22using firefox binary path from environment variable %25s: %25s%25n%22, ENV_FIREFOX_BIN, firefoxBin);%0A                        return WebDriverSupport.firefoxOnDisplay(display).create(new FirefoxBinary(new File(firefoxBin)), new FirefoxProfile());%0A                    %7D else %7B%0A    
                 
@@ -5398,32 +5398,54 @@
 play).create();%0A
+                    %7D%0A
                 

support FIREFOX_BIN env var for firefox binary specification