package com.codeborne.selenide.impl;

import com.codeborne.selenide.Config;
import org.openqa.selenium.Alert;
import org.openqa.selenium.UnhandledAlertException;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebDriverException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.CheckReturnValue;
import javax.annotation.Nonnull;
import javax.annotation.ParametersAreNonnullByDefault;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;

import static java.nio.charset.StandardCharsets.UTF_8;

@ParametersAreNonnullByDefault
public class WebPageSourceExtractor implements PageSourceExtractor {
  private static final Logger log = LoggerFactory.getLogger(WebPageSourceExtractor.class);
  protected Set<String> printedErrors = new ConcurrentSkipListSet<>();

  @Nonnull
  @CheckReturnValue
  @Override
  public File extract(Config config, WebDriver driver, String fileName) {
    return extract(config, driver, fileName, true);
  }

  private File extract(Config config, WebDriver driver, String fileName, boolean retryIfAlert) {
    File pageSource = new File(config.reportsFolder(), fileName + ".html");
    try {
      writeToFile(driver.getPageSource(), pageSource);
    } catch (UnhandledAlertException e) {
      if (retryIfAlert) {
        retryingExtractionOnAlert(config, driver, fileName, e);
      } else {
        printOnce("savePageSourceToFile", e);
      }
    } catch (WebDriverException e) {
      log.warn("Failed to save page source to {}", fileName, e);
      writeToFile(e.toString(), pageSource);
      return pageSource;
    } catch (RuntimeException e) {
      log.error("Failed to save page source to {}", fileName, e);
      writeToFile(e.toString(), pageSource);
    }
    return pageSource;
  }

  protected void writeToFile(String content, File targetFile) {
    try (ByteArrayInputStream in = new ByteArrayInputStream(content.getBytes(UTF_8))) {
      FileHelper.copyFile(in, targetFile);
    } catch (IOException e) {
      log.error("Failed to write file {}", targetFile.getAbsolutePath(), e);
    }
  }

  protected synchronized void printOnce(String action, Throwable error) {
    if (!printedErrors.contains(action)) {
      log.error(error.getMessage(), error);
      printedErrors.add(action);
    } else {
      log.error("Failed to {}: {}", action, error);
    }
  }

  private void retryingExtractionOnAlert(Config config, WebDriver driver, String fileName, Exception e) {
    try {
      Alert alert = driver.switchTo().alert();
      log.error("{}: {}", e, alert.getText());
      alert.accept();
      extract(config, driver, fileName, false);
    } catch (Exception unableToCloseAlert) {
      log.error("Failed to close alert", unableToCloseAlert);
    }
  }
}

@@ -839,23 +839,27 @@
 s);%0A  pr
-otected
+ivate final
  Set%3CStr
@@ -1204,20 +1204,22 @@
 ource = 
-new 
+create
 File(con
@@ -1225,44 +1225,18 @@
 nfig
-.reportsFolder(), fileName + %22.html%22
+, fileName
 );%0A 
@@ -1296,32 +1296,36 @@
 geSource);%0A    %7D
+%0A   
  catch (Unhandle
@@ -1438,24 +1438,30 @@
  e);%0A      %7D
+%0A     
  else %7B%0A    
@@ -1511,24 +1511,28 @@
      %7D%0A    %7D
+%0A   
  catch (WebD
@@ -1687,24 +1687,28 @@
 ource;%0A    %7D
+%0A   
  catch (Runt
@@ -1866,24 +1866,167 @@
 ource;%0A  %7D%0A%0A
+  @Nonnull%0A  protected File createFile(Config config, String fileName) %7B%0A    return new File(config.reportsFolder(), fileName + %22.html%22);%0A  %7D%0A%0A
   protected 
@@ -2205,32 +2205,36 @@
 rgetFile);%0A    %7D
+%0A   
  catch (IOExcept
@@ -2529,16 +2529,20 @@
 );%0A    %7D
+%0A   
  else %7B%0A
@@ -2884,21 +2884,25 @@
 false);%0A
-
     %7D
+%0A   
  catch (

make WebPageSourceExtractor easily overridable