package org.infinispan.commons.test;


import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamWriter;

/**
 * A JUnit XML report generator for Polarion based on the JUnitXMLReporter
 *
 * <p>Extracted from {@link PolarionJUnitXMLReporter}</p>
 *
 * @author <a href='mailto:afield[at]redhat[dot]com'>Alan Field</a>
 * @author Dan Berindei
 * @since 10.0
 */
public class PolarionJUnitXMLWriter implements AutoCloseable {

   private final FileWriter fileWriter;

   public enum Status {SUCCESS, FAILURE, ERROR, SKIPPED}

   private static final String TESTSUITE = "testsuite";
   private static final String ATTR_TESTS = "tests";
   private static final String ATTR_TIME = "time";
   private static final String ATTR_NAME = "name";
   private static final String ATTR_SKIPPED = "skipped";
   private static final String ATTR_ERRORS = "errors";
   private static final String ATTR_FAILURES = "failures";

   private static final String TESTCASE = "testcase";
   private static final String FAILURE = "failure";
   private static final String RERUN_FAILURE = "rerunFailure";
   private static final String FLAKY_FAILURE = "flakyFailure";
   private static final String ERROR = "error";
   private static final String SKIPPED = "skipped";
   private static final String ATTR_CLASSNAME = "classname";
   private static final String ATTR_MESSAGE = "message";
   private static final String ATTR_TYPE = "type";
   private static final String ATTR_VALUE = "value";
   private static final String STACKTRACE = "stackTrace";

   private static final String PROPERTIES = "properties";
   private static final String PROPERTY = "property";

   private static final char UNICODE_REPLACEMENT = 0xFFFD;

   private static final XMLOutputFactory xmlOutputFactory = XMLOutputFactory.newFactory();

   private XMLStreamWriter xmlWriter;

   public PolarionJUnitXMLWriter(File outputFile) throws IOException {
      File outputDirectory = outputFile.getParentFile();
      if (!outputDirectory.exists()) {
         if (!outputDirectory.mkdirs()) {
            throw new IllegalStateException("Unable to create report directory " + outputDirectory);
         }
      }
      fileWriter = new FileWriter(outputFile);
   }

   public void start(String testsuiteName, long testCount, long skippedCount, long failedCount, long elapsedTime,
                     boolean includeProperties) throws XMLStreamException {
      xmlWriter = new PrettyXMLStreamWriter(xmlOutputFactory.createXMLStreamWriter(fileWriter));

      xmlWriter.writeStartDocument();
      xmlWriter.writeCharacters("\n");
      xmlWriter.writeComment("Generated by " + getClass().getName());

      xmlWriter.writeStartElement(TESTSUITE);
      xmlWriter.writeAttribute(ATTR_TESTS, "" + testCount);
      xmlWriter.writeAttribute(ATTR_TIME, "" + elapsedTime / 1000.0);
      xmlWriter.writeAttribute(ATTR_NAME, testsuiteName);
      xmlWriter.writeAttribute(ATTR_SKIPPED, "" + skippedCount);
      xmlWriter.writeAttribute(ATTR_ERRORS, "0");
      xmlWriter.writeAttribute(ATTR_FAILURES, "" + failedCount);

      if (includeProperties) {
         writeProperties();
      }

      xmlWriter.writeComment("Tests results");
   }

   @Override
   public void close() throws Exception {
      xmlWriter.writeEndElement();
      xmlWriter.writeEndDocument();
      xmlWriter.close();
      fileWriter.close();
   }

   public void writeTestCase(PolarionJUnitTest test) throws XMLStreamException {
      PolarionJUnitTest.Status status = test.status;
      if (status == PolarionJUnitTest.Status.SUCCESS) {
         xmlWriter.writeEmptyElement(TESTCASE);
         writeTestAttributes(test);
         return;
      }
      xmlWriter.writeStartElement(TESTCASE);
      writeTestAttributes(test);
      switch (test.status) {
         case FLAKY:
            writeCauseElements(FLAKY_FAILURE, test.failures);
            break;
         case ERROR:
            writeCauseElements(ERROR, test.failures);
            break;
         case FAILURE:
            writeCauseElement(FAILURE, test.failures.get(0));
            writeCauseElements(RERUN_FAILURE, test.failures.subList(1, test.failures.size()));
            break;
         case SKIPPED:
            xmlWriter.writeEmptyElement(SKIPPED);
      }
      xmlWriter.writeEndElement();
   }

   public void writeTestAttributes(PolarionJUnitTest test) throws XMLStreamException {
      xmlWriter.writeAttribute(ATTR_NAME, test.name);
      xmlWriter.writeAttribute(ATTR_CLASSNAME, test.clazz);
      xmlWriter.writeAttribute(ATTR_TIME, Double.toString(test.elapsedTime() / 1000.0));
   }

   private void writeCauseElements(String tag, Collection<Throwable> throwables) throws XMLStreamException {
      for (Throwable t : throwables)
         writeCauseElement(tag, t);
   }

   private void writeCauseElement(String tag, Throwable throwable) throws XMLStreamException {
      String throwableClass = throwable.getClass().getName();
      String message = throwable.getMessage();
      StringWriter sw = new StringWriter();
      PrintWriter pw = new PrintWriter(sw);
      throwable.printStackTrace(pw);
      pw.flush();
      String stackTrace = escapeHtml(sw.toString());

      xmlWriter.writeStartElement(tag);
      xmlWriter.writeAttribute(ATTR_TYPE, throwableClass);
      if ((message != null) && !message.isEmpty()) {
         xmlWriter.writeAttribute(ATTR_MESSAGE, escapeInvalidChars(message));
      }
      xmlWriter.writeStartElement(STACKTRACE);
      xmlWriter.writeCData(stackTrace);
      xmlWriter.writeEndElement();
      xmlWriter.writeEndElement();
   }

   private String escapeInvalidChars(String chars) {
      // Restricted chars for xml are [#x1-#x8], [#xB-#xC], [#xE-#x1F], [#x7F-#x84], [#x86-#x9F]
      return chars.codePoints()
                  .map(c -> {
                     if (!Character.isDefined(c) || 0x1 <= c && c <= 0x8 || c == 0xB || c == 0xC ||
                         0xE <= c && c <= 0x1F || 0x7F <= c && c < 0x84 || 0x86 <= c && c <= 0x9F) {
                        return UNICODE_REPLACEMENT;
                     } else {
                        return c;
                     }
                  })
                  .collect(StringBuilder::new, StringBuilder::appendCodePoint, StringBuilder::append)
                  .toString();
   }

   private void writeProperties() throws XMLStreamException {
      xmlWriter.writeStartElement(PROPERTIES);

      // Add all system properties
      xmlWriter.writeComment("Java System properties");
      for (Object key : System.getProperties().keySet()) {
         xmlWriter.writeEmptyElement(PROPERTY);
         xmlWriter.writeAttribute(ATTR_NAME, key.toString());
         xmlWriter.writeAttribute(ATTR_VALUE, System.getProperty(key.toString()));
      }

      // Add all environment variables
      xmlWriter.writeComment("Environment variables");
      for (String key : System.getenv().keySet()) {
         xmlWriter.writeEmptyElement(PROPERTY);
         xmlWriter.writeAttribute(ATTR_NAME, key);
         xmlWriter.writeAttribute(ATTR_VALUE, System.getenv(key));
      }

      xmlWriter.writeEndElement();
   }

   private static final Map<Character, String> ESCAPES = new HashMap<>();

   static {
      ESCAPES.put('<', "&lt;");
      ESCAPES.put('>', "&gt;");
      ESCAPES.put('\'', "&apos;");
      ESCAPES.put('"', "&quot;");
      ESCAPES.put('&', "&amp;");
   }

   public static String escapeHtml(String s) {
      if (s == null) {
         return null;
      }
      StringBuilder result = new StringBuilder();
      for (int i = 0; i < s.length(); i++) {
         char c = s.charAt(i);
         String nc = ESCAPES.get(c);
         if (nc != null) {
            result.append(nc);
         } else {
            result.append(c);
         }
      }
      return result.toString();
   }
}
