package io.groundhog.jmeter;

import io.groundhog.base.URIScheme;
import io.groundhog.replay.ReplayClient;

import com.google.common.net.HostAndPort;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Module;
import org.apache.jmeter.samplers.AbstractSampler;
import org.apache.jmeter.samplers.Entry;
import org.apache.jmeter.samplers.Interruptible;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.testbeans.TestBean;
import org.apache.jmeter.testelement.ThreadListener;
import org.apache.jmeter.threads.JMeterContextService;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;

import java.io.File;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

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

/**
 * @author Danny Thomas
 * @since 1.0
 */
public class HarReplaySampler extends AbstractSampler implements TestBean, ThreadListener, Interruptible {
  private static final Logger LOG = LoggingManager.getLoggerForClass();

  private final BlockingQueue<SampleResult> results = new LinkedBlockingQueue<>();

  private String filename;
  private URIScheme scheme;
  private String host;
  private int port;
  private ReplayClient client;

  @Override
  public SampleResult sample(Entry entry) {
    if (null == client) {
      HostAndPort hostAndPort = HostAndPort.fromParts(host, port);
      Module jmeterModule = new JMeterModule(new File(filename), results, scheme, hostAndPort);
      Injector injector = Guice.createInjector(new JMeterSlf4jModule(), jmeterModule);
      client = injector.getInstance(ReplayClient.class);

      LOG.info("Starting replay for " + filename + " using " + scheme.getScheme() + " against " + hostAndPort);
      client.startAsync();
      client.awaitRunning();
    }

    while (true) {
      if (client.isRunning() || !results.isEmpty()) {
        try {
          SampleResult result = results.poll(500, TimeUnit.MILLISECONDS);
          if (null != result) {
            return result;
          }
        } catch (InterruptedException e) {
          break;
        }
      } else {
        LOG.info("Replay client has stopped. Stopping current thread");
        JMeterContextService.getContext().getThread().stop();
        break;
      }
    }
    return null;
  }

  @Override
  public boolean interrupt() {
    threadFinished();
    return true;
  }

  @Override
  public void threadStarted() {
  }

  @Override
  public void threadFinished() {
    if (null != client && client.isRunning()) {
      client.stopAsync();
      client.awaitTerminated();
    }
  }

  public String getFilename() {
    return filename;
  }

  public void setFilename(String filename) {
    this.filename = checkNotNull(filename);
  }

  public int getScheme() {
    return scheme.ordinal();
  }

  public void setScheme(int scheme) {
    this.scheme = URIScheme.values()[scheme];
  }

  public String getHost() {
    return host;
  }

  public void setHost(String host) {
    this.host = checkNotNull(host);
  }

  public int getPort() {
    return port;
  }

  public void setPort(int port) {
    this.port = port;
  }
}

@@ -99,16 +99,62 @@
 yClient;
+%0Aimport io.groundhog.replay.RequestDispatcher;
 %0A%0Aimport
@@ -1448,24 +1448,38 @@
 = client) %7B%0A
+      try %7B%0A  
       HostAn
@@ -1539,16 +1539,18 @@
 ;%0A      
+  
 Module j
@@ -1627,24 +1627,26 @@
 stAndPort);%0A
+  
       Inject
@@ -1722,24 +1722,26 @@
 ule);%0A      
+  
 client = inj
@@ -1782,24 +1782,26 @@
 ss);%0A%0A      
+ 
+ 
 LOG.info(%22St
@@ -1896,24 +1896,26 @@
 ort);%0A      
+  
 client.start
@@ -1915,32 +1915,34 @@
 t.startAsync();%0A
+  
       client.awa
@@ -1950,24 +1950,212 @@
 tRunning();%0A
+      %7D catch (Exception e) %7B%0A        LOG.error(%22Fatal error while initialising replay. Stopping current thread%22, e);%0A        JMeterContextService.getContext().getThread().stop();%0A      %7D%0A
     %7D%0A%0A    w
@@ -2285,11 +2285,47 @@
 oll(
-500
+RequestDispatcher.SKEW_THRESHOLD_MILLIS
 , Ti

Improve behaviour of sampler when a fatal error occurs during replay initialisation