/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 com.webautomation;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URL;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import org.apache.commons.io.FilenameUtils;
import org.apache.commons.io.IOUtils;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.htmlunit.HtmlUnitDriver;
import org.openqa.selenium.internal.Base64Encoder;
import org.openqa.selenium.remote.CapabilityType;
import org.openqa.selenium.remote.DesiredCapabilities;

import com.gargoylesoftware.htmlunit.BrowserVersion;
import com.gargoylesoftware.htmlunit.WebClient;
import com.gargoylesoftware.htmlunit.WebRequest;
import com.gargoylesoftware.htmlunit.WebWindow;
import com.gargoylesoftware.htmlunit.html.HtmlElement;
import com.gargoylesoftware.htmlunit.html.HtmlPage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * from https://code.google.com/p/selenium/issues/detail?id=1361
 */
public class ScreenCaptureHtmlUnitDriver extends HtmlUnitDriver implements TakesScreenshot {

    private static Map<String, byte[]> imagesCache = Collections.synchronizedMap(new HashMap<String, byte[]>());

    private static Map<String, String> cssjsCache = Collections.synchronizedMap(new HashMap<String, String>());

    // http://stackoverflow.com/questions/4652777/java-regex-to-get-the-urls-from-css
    private final static Pattern cssUrlPattern = Pattern.compile("background(-image)?[\\s]*:[^url]*url[\\s]*\\([\\s]*([^\\)]*)[\\s]*\\)[\\s]*");// ?<url>

    static Logger LOGGER = LoggerFactory.getLogger(ScreenCaptureHtmlUnitDriver.class);

    public ScreenCaptureHtmlUnitDriver() {
        super();
    }

    public ScreenCaptureHtmlUnitDriver(boolean enableJavascript) {
        super(enableJavascript);
    }

    public ScreenCaptureHtmlUnitDriver(Capabilities capabilities) {
        super(capabilities);
    }

    public ScreenCaptureHtmlUnitDriver(BrowserVersion version) {
        super(version);
        DesiredCapabilities var = ((DesiredCapabilities) getCapabilities());
        var.setCapability(CapabilityType.TAKES_SCREENSHOT, true);
    }

    @Override
    @SuppressWarnings("unchecked")
    public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
        byte[] archive = new byte[0];
        try {
            archive = downloadCssAndImages(getWebClient(), (HtmlPage) getCurrentWindow().getEnclosedPage());
        } catch (Exception e) {
            LOGGER.error("Exception in ScreenCaptureHtmlUnitDriver while getScreenshotAs ", e);
        }
        if(target.equals(OutputType.BASE64)){
            return target.convertFromBase64Png(new Base64Encoder().encode(archive));
        }
        if(target.equals(OutputType.FILE)){
            File f = new File("screen.tmp");
            try {
                FileOutputStream scr = new FileOutputStream(f);
                scr.write(archive);
                scr.close();
            } catch (IOException e) {
                throw new WebDriverException(e);
            }
            return (X) f;
        }
        return (X) archive;
    }

    // http://stackoverflow.com/questions/2244272/how-can-i-tell-htmlunits-webclient-to-download-images-and-css
    protected byte[] downloadCssAndImages(WebClient webClient, HtmlPage page) throws Exception {
        WebWindow currentWindow = webClient.getCurrentWindow();
        Map<String, String> urlMapping = new HashMap<String, String>();
        Map<String, byte[]> files = new HashMap<String, byte[]>();
        WebWindow window = null;
        try {
            window = webClient.getWebWindowByName(page.getUrl().toString()+"_screenshot");
            webClient.getPage(window, new WebRequest(page.getUrl()));
        } catch (Exception e) {
            LOGGER.error("Exception in ScreenCaptureHtmlUnitDriver while downloadCssAndImages ", e);
            window = webClient.openWindow(page.getUrl(), page.getUrl().toString()+"_screenshot");
        }

        String xPathExpression = "//*[name() = 'img' or name() = 'link' and (@type = 'text/css' or @type = 'image/x-icon') or  @type = 'text/javascript']";
        List<?> resultList = page.getByXPath(xPathExpression);

        Iterator<?> i = resultList.iterator();
        while (i.hasNext()) {
            try {
                HtmlElement el = (HtmlElement) i.next();
                String resourceSourcePath = el.getAttribute("src").equals("") ? el.getAttribute("href") : el
                        .getAttribute("src");
                if (resourceSourcePath == null || resourceSourcePath.equals(""))
                    continue;
                URL resourceRemoteLink = page.getFullyQualifiedUrl(resourceSourcePath);
                String resourceLocalPath = mapLocalUrl(page, resourceRemoteLink, resourceSourcePath, urlMapping);
                urlMapping.put(resourceSourcePath, resourceLocalPath);
                if (!resourceRemoteLink.toString().endsWith(".css")) {
                    byte[] image = downloadImage(webClient, window, resourceRemoteLink);
                    files.put(resourceLocalPath, image);
                } else {
                    String css = downloadCss(webClient, window, resourceRemoteLink);
                    for (String cssImagePath : getLinksFromCss(css)) {
                        URL cssImagelink = page.getFullyQualifiedUrl(cssImagePath.replace("\"", "").replace("\'", "")
                                .replace(" ", ""));
                        String cssImageLocalPath = mapLocalUrl(page, cssImagelink, cssImagePath, urlMapping);
                        files.put(cssImageLocalPath, downloadImage(webClient, window, cssImagelink));
                    }
                    files.put(resourceLocalPath, replaceRemoteUrlsWithLocal(css, urlMapping)
                            .replace("resources/", "./").getBytes());
                }
            } catch (Exception e) {
                LOGGER.error("Exception in ScreenCaptureHtmlUnitDriver while resultList.iterator ", e);
            }
        }
        String pagesrc = replaceRemoteUrlsWithLocal(page.getWebResponse().getContentAsString(), urlMapping);
        files.put("page.html", pagesrc.getBytes());
        webClient.setCurrentWindow(currentWindow);
        return createZip(files);
    }

    String downloadCss(WebClient webClient, WebWindow window, URL resourceUrl) throws Exception {
        if (cssjsCache.get(resourceUrl.toString()) == null) {
            cssjsCache.put(resourceUrl.toString(), webClient.getPage(window, new WebRequest(resourceUrl))
                    .getWebResponse().getContentAsString());

        }
        return cssjsCache.get(resourceUrl.toString());
    }

    byte[] downloadImage(WebClient webClient, WebWindow window, URL resourceUrl) throws Exception {
        if (imagesCache.get(resourceUrl.toString()) == null) {
            imagesCache.put(
                    resourceUrl.toString(),
                    IOUtils.toByteArray(webClient.getPage(window, new WebRequest(resourceUrl)).getWebResponse()
                            .getContentAsStream()));
        }
        return imagesCache.get(resourceUrl.toString());
    }

    public static byte[] createZip(Map<String, byte[]> files) throws IOException {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        ZipOutputStream zipfile = new ZipOutputStream(bos);
        Iterator<String> i = files.keySet().iterator();
        String fileName = null;
        ZipEntry zipentry = null;
        while (i.hasNext()) {
            fileName = i.next();
            zipentry = new ZipEntry(fileName);
            zipfile.putNextEntry(zipentry);
            zipfile.write(files.get(fileName));
        }
        zipfile.close();
        return bos.toByteArray();
    }

    List<String> getLinksFromCss(String css) {
        List<String> result = new LinkedList<String>();
        Matcher m = cssUrlPattern.matcher(css);
        while (m.find()) { // find next match
            result.add( m.group(2));
        }
        return result;
    }

    String replaceRemoteUrlsWithLocal(String source, Map<String, String> replacement) {
        for (String object : replacement.keySet()) {
            // background:url(http://org.com/images/image.gif)
            source = source.replace(object, replacement.get(object));
        }
        return source;
    }

    String mapLocalUrl(HtmlPage page, URL link, String path, Map<String, String> replacementToAdd) throws Exception {
        String resultingFileName = "resources/" + FilenameUtils.getName(link.getFile());
        replacementToAdd.put(path, resultingFileName);
        return resultingFileName;
    }

}