package com.carrotsearch.ant.tasks.junit4.listeners;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

import org.apache.tools.ant.BuildException;
import org.apache.tools.ant.Project;
import org.apache.tools.ant.types.Resource;
import org.apache.tools.ant.types.ResourceCollection;

import com.carrotsearch.ant.tasks.junit4.JUnit4;
import com.carrotsearch.ant.tasks.junit4.events.aggregated.AggregatedQuitEvent;
import com.carrotsearch.ant.tasks.junit4.events.aggregated.AggregatedSuiteResultEvent;
import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.eventbus.Subscribe;
import com.google.common.io.Closer;
import com.google.common.io.Files;

/**
 * A report listener that emits per-suite execution times information useful
 * for load balancing tests across JVMs. 
 */
public class ExecutionTimesReport implements AggregatedEventListener {
  /**
   * @see #setHistoryLength(int)
   */
  public final static int DEFAULT_HISTORY_LENGTH = 10;

  /**
   * The file where suite hints are stored/ updated.
   */
  private File hintsFile;

  /**
   * Execution time hints. Key: suite name, value: execution times.
   */
  private Map<String, List<Long>> hints; 

  /**
   * @see #setHistoryLength(int)
   */
  private int historyLength = DEFAULT_HISTORY_LENGTH;

  /**
   * Outer task.
   */
  private JUnit4 outer;
  
  /**
   * Hints file (key-value pairs).
   */
  public void setFile(File hintsFile) {
    this.hintsFile = hintsFile;
  }
  
  /**
   * How many execution times to store per-suite? The history must be larger than 0.
   */
  public void setHistoryLength(int length) {
    if (length < 0) {
      throw new BuildException("History length must be >= 1: " + length);
    }
    this.historyLength = length;
  }

  /**
   * Remember execution time for all executed suites.
   */
  @Subscribe
  public void onSuiteResult(AggregatedSuiteResultEvent e) {
    long millis = e.getExecutionTime();
    String suiteName = e.getDescription().getDisplayName();
    
    List<Long> values = hints.get(suiteName);
    if (values == null) {
      hints.put(suiteName, values = Lists.newArrayList());
    }
    values.add(millis);
    while (values.size() > historyLength)
      values.remove(0);
  }
  
  /**
   * Write back to hints file.
   */
  @Subscribe
  public void onEnd(AggregatedQuitEvent e) {
    try {
      writeHints(hintsFile, hints);
    } catch (IOException exception) {
      outer.log("Could not write back the hints file.", exception, Project.MSG_ERR);
    }
  }

  @Override
  public void setOuter(JUnit4 outer) throws BuildException {
    if (hintsFile == null) {
      throw new BuildException(
          "Execution times listener requires file attribute.");
    }
    
    // If the file already exists, read its contents.
    try {
      if (hintsFile.isFile()) {
          hints = readHints(hintsFile);
      } else {
          if (!hintsFile.createNewFile()) {
            throw new BuildException("Could not create file: "
                + hintsFile.getAbsolutePath());
          }
          hints = Maps.newHashMap();
      }
    } catch (IOException e) {
      throw new BuildException("Could not read or create hints file: " + hintsFile.getAbsolutePath(), e);
    }
    
    this.outer = outer;
  }

  /**
   * Read hints from a file.
   */
  public static Map<String,List<Long>> readHints(File hints) throws IOException {
    Map<String,List<Long>> result = Maps.newHashMap();
    InputStream is = new FileInputStream(hints);
    mergeHints(is, result);
    return result;
  }

  /**
   * Read hints from a file and merge with the given hints map.
   */
  public static void mergeHints(InputStream is, Map<String,List<Long>> hints) throws IOException {
    final BufferedReader reader = new BufferedReader(
        new InputStreamReader(is, Charsets.UTF_8));
    
    String line;
    while ((line = reader.readLine()) != null) {
      line = line.trim();
      if (line.isEmpty() || line.startsWith("#"))
        continue;

      final int equals = line.indexOf('=');
      if (equals <= 0) {
        throw new IOException("No '=' character on a non-comment line?: " + line);
      } else {
        String key = line.substring(0, equals);
        List<Long> values = hints.get(key);
        if (values == null) {
          hints.put(key, values = Lists.newArrayList());
        }
        for (String v : line.substring(equals + 1).split("[\\,]")) {
          if (!v.isEmpty()) values.add(Long.parseLong(v));
        }
      }
    }
  }

  /**
   * Writes back hints file. 
   */
  public static void writeHints(File file, Map<String,List<Long>> hints) throws IOException {
    Closer closer = Closer.create();
    try {
      BufferedWriter w = closer.register(Files.newWriter(file, Charsets.UTF_8));
      if (!(hints instanceof SortedMap)) {
        hints = new TreeMap<String,List<Long>>(hints);
      }
      
      Joiner joiner = Joiner.on(',');
      for (Map.Entry<String,List<Long>> e : hints.entrySet()) {
        w.write(e.getKey());
        w.write("=");
        joiner.appendTo(w, e.getValue());
        w.write("\n");
      }
    } catch (Throwable t) {
      throw closer.rethrow(t);
    } finally {
      closer.close();
    }
  }

  /**
   * Read hints from all resources in a collection, retaining 
   * <code>suiteNames</code>. If <code>suiteNames</code> is null,
   * everything is retained.
   */
  public static Map<String,List<Long>> mergeHints(
      Collection<ResourceCollection> resources, Collection<String> suiteNames) {
    final Map<String,List<Long>> hints = Maps.newHashMap();
    for (ResourceCollection rc : resources) {
      @SuppressWarnings("unchecked")
      Iterator<Resource> i = rc.iterator();
      while (i.hasNext()) {
        InputStream is = null;
        Resource r = i.next();
        try {
          is = r.getInputStream();
          mergeHints(is, hints);

          // Early prune the hints to those we have on the list.
          if (suiteNames != null) {
            hints.keySet().retainAll(suiteNames);
          }
        } catch (IOException e) {
          throw new BuildException("Could not read hints from resource: "
              + r.getDescription(), e);
        } finally {
          try {
            if (is != null) is.close();
          } catch (IOException e) {
            throw new BuildException("Could not close hints file: " + r.getDescription());
          }
        }
      }
    }
    return hints;
  }  
}
