var _ = require('underscore'),
    context = require('./context'),
    Mocha = require('mocha'),
    Spooky = require('./spooky'),
    temp = require('temp');

module.exports.config = require('./config');
module.exports.context = require('./context');

// If we are run by global mocha, sniff for that case.
process.mainModule.children.forEach(function(child) {
  if (/mocha[\/\\]index.js/.test(child.filename)) {
    Mocha = child.exports;
  }
});

var loadFiles = Mocha.prototype.loadFiles;
Mocha.prototype.loadFiles = function() {
  if (!this.suite._carahueInit) {
    this.suite.beforeEach(function(done) {
      if (!Spooky.isActive()) {
        Spooky.get(function() {
          // Need this indirection to prevent param passing to mocha
          done();
        });
      } else {
        done();
      }
    });

    this.suite.beforeAll(function(done) {
      var self = this;

      temp.mkdir('carahue', function(err, dir) {
        if (err) {
          throw err;
        }
        self.tmpDir = dir;
        done();
      });
    });

    this.suite.on('pre-require', function(context) {
      context.beforePage = beforePage;
      context.page = page;
    });

    this.suite._carahueInit = true;
  }

  return loadFiles.apply(this);
};


function beforePage(fn) {
  before(function() {
    this.beforePage = this.beforePage || [];
    this.beforePage.push(fn);
  });
  after(function() {
    var index = _.lastIndexOf(this.beforePage, fn);
    if (index >= 0) {
      this.beforePage.splice(index, 1);
    }
  });
}

function page(name, page, fn) {
  it(name, function(done) {
    var self = this;

    start = Date.now();

    Spooky.get(function(spooky) {
      context.inject(self, spooky);

      if (self.routePrefix && !(/^\//.test(page) || /:\/\//.test(page))) {
        page = self.routePrefix + page;
      }

      function cleanup() {
        spooky.removeListener('carahue.done', onDone);
        spooky.removeListener('carahue.capture', onCapture);
        spooky.removeListener('carahue.capture.complete', onCaptureComplete);
      }

      function onDone(isDone) {
        self.done = self.done || isDone;

        if (!self.screencaptureCount && self.done) {
          cleanup();
          done();
        }
      }

      self.screencaptureCount = 0;
      self.done = false;
      function onCapture() {
        self.screencaptureCount++;
      }
      function onCaptureComplete() {
        self.screencaptureCount--;
        this.emit('carahue.done');
      }

      spooky.on('carahue.done', onDone);
      spooky.on('carahue.capture', onCapture);
      spooky.on('carahue.capture.complete', onCaptureComplete);

      spooky.start(page);
      _.each(self.beforePage, function(before) {
        before.call(self, name, page);
      });

      self.thenScreenshot('');

      try {
        fn && fn.call(self);
      } catch (err) {
        cleanup();

        throw err;
      }

      spooky.run(function() {
        this.emit('carahue.done', true);
      });
    });
  });
}

@@ -484,16 +484,40 @@
 oadFiles
+,%0A    cleanupEvents = %7B%7D
 ;%0AMocha.
@@ -831,32 +831,431 @@
      %7D%0A    %7D);%0A%0A
+    // Event listener cleanup%0A    this.suite.afterEach(function(done) %7B%0A      if (Spooky.isActive()) %7B%0A        Spooky.get(function(spooky) %7B%0A          _.each(cleanupEvents, function(handler, name) %7B%0A            spooky.removeListener(name, handler);%0A          %7D);%0A          cleanupEvents = %7B%7D;%0A          done();%0A        %7D);%0A      %7D else %7B%0A        cleanupEvents = %7B%7D;%0A        done();%0A      %7D%0A    %7D);%0A%0A
     this.suite.b
@@ -2087,24 +2087,132 @@
 n(spooky) %7B%0A
+      function register(name, fn) %7B%0A        spooky.on(name, fn);%0A        cleanupEvents%5Bname%5D = fn;%0A      %7D%0A%0A
       contex
@@ -2370,57 +2370,15 @@
     
-function cleanup() %7B%0A        spooky.removeL
+reg
 iste
-ne
 r('c
@@ -2394,171 +2394,8 @@
 ne',
- onDone);%0A        spooky.removeListener('carahue.capture', onCapture);%0A        spooky.removeListener('carahue.capture.complete', onCaptureComplete);%0A      %7D%0A%0A     
  fun
@@ -2399,23 +2399,16 @@
 function
- onDone
 (isDone)
@@ -2509,29 +2509,8 @@
 ) %7B%0A
-          cleanup();%0A
     
@@ -2532,32 +2532,34 @@
        %7D%0A      %7D
+);
 %0A%0A      self.scr
@@ -2610,34 +2610,52 @@
 ;%0A      
-function onCapture
+register('carahue.capture', function
 () %7B%0A   
@@ -2697,274 +2697,198 @@
    %7D
+);
 %0A      
-function onC
+register('carahue.c
 apture
-C
+.c
 omplete
-() %7B%0A        self.screencaptureCount--;%0A        this.emit('carahue.done');%0A      %7D%0A%0A      spooky.on('carahue.done', onDone);%0A      spooky.on('carahue.capture', onCapture);%0A      spooky.on('carahue.capture.complete', onCaptureComplete
+', function(err) %7B%0A        if (err) %7B%0A          throw err;%0A        %7D%0A%0A        self.screencaptureCount--;%0A        this.emit('carahue.done');%0A      %7D
 );%0A%0A
@@ -3048,22 +3048,8 @@
 );%0A%0A
-      try %7B%0A  
     
@@ -3074,77 +3074,8 @@
 lf);
-%0A      %7D catch (err) %7B%0A        cleanup();%0A%0A        throw err;%0A      %7D
 %0A%0A  

Move spooky event cleanup to afterEach step