import expect from 'unexpected';
import avaTest from 'ava/lib/test';
import { AssertionError } from 'ava/lib/assert';

const ec = Object.getPrototypeOf(avaTest.prototype.createExecutionContext());

ec.expect = function __$AVA_UNEXPECTED_EXPECT$__(...args) {
  const onFail = origErr => {
    let { stack } = origErr;
    const cutoff = stack.indexOf('__$AVA_UNEXPECTED_EXPECT$__');
    if (cutoff >= 0) {
      stack = stack.substr(0, stack.lastIndexOf('\n', cutoff))
        + stack.substr(stack.indexOf('\n', cutoff));
    }
    const err = new AssertionError({
      assertion: 'fail',
      message: origErr.message,
      stack,
    });
    this._test.addFailedAssertion(err);
  };
  let expectation;
  try {
    expectation = expect(...args);
  } catch (origErr) {
    onFail(origErr);
    return;
  }
  if (expectation.isFulfilled()) {
    if (expectation.isRejected()) {
      onFail(expectation.reason());
    } else {
      this.pass();
    }
  } else {
    return expectation.then(() => this.pass(), onFail);
  }
};
