#include <ruby.h>
#include <env.h>
#include <node.h>
#include <st.h>
#include <stdlib.h>
#include <assert.h>

#define COVERAGE_DEBUG_EVENTS 0
#define RCOVRT_VERSION_MAJOR 2
#define RCOVRT_VERSION_MINOR 0
#define RCOVRT_VERSION_REV   0

static VALUE mRcov;
static VALUE mRCOV__;
static VALUE oSCRIPT_LINES__;
static ID id_cover;
static st_table* coverinfo = 0;
static char coverage_hook_set_p;

struct cov_array {
  unsigned int len;
  unsigned int *ptr;
};

static struct cov_array *cached_array = 0;
static char *cached_file = 0;

static struct cov_array * coverage_increase_counter_uncached(char *sourcefile, unsigned int sourceline, char mark_only) {
  struct cov_array *carray = NULL;

  if (sourcefile == NULL) {
    return NULL;
  } else if (!st_lookup(coverinfo, (st_data_t)sourcefile, (st_data_t*)&carray)) {
    VALUE arr;
    arr = rb_hash_aref(oSCRIPT_LINES__, rb_str_new2(sourcefile));

    if (NIL_P(arr)) return 0;

    rb_check_type(arr, T_ARRAY);
    carray = calloc(1, sizeof(struct cov_array));
    carray->ptr = calloc(RARRAY(arr)->len, sizeof(unsigned int));
    carray->len = RARRAY(arr)->len;
    st_insert(coverinfo, (st_data_t)strdup(sourcefile), (st_data_t) carray);
  } else {
    assert(carray && "failed to create valid carray");
  }

  if (mark_only) {
    if (carray && carray->len > sourceline) {
      if (!carray->ptr[sourceline])
        carray->ptr[sourceline] = 1;
    } else {
#if COVERAGE_DEBUG_EVENTS
      printf("DEBUG: %s carray->len:%d sourceline:%d\n", sourcefile, carray->len, sourceline);
#endif
    }
  } else {
    if (carray && carray->len > sourceline) {
      carray->ptr[sourceline]++;
    }
  }

  return carray;
}

static void coverage_mark_caller() {
  struct FRAME *frame = ruby_frame;
  NODE *n;

  if (frame->last_func == ID_ALLOCATOR) frame = frame->prev;

  for (; frame && (n = frame->node); frame = frame->prev) {
    if (frame->prev && frame->prev->last_func) {
      if (frame->prev->node == n) {
        if (frame->prev->last_func == frame->last_func) continue;
      }
      coverage_increase_counter_uncached(n->nd_file, nd_line(n) - 1, 1);
    } else {
      coverage_increase_counter_uncached(n->nd_file, nd_line(n) - 1, 1);
    }
    break;
  }
}

static void coverage_increase_counter_cached(char *sourcefile, int sourceline) {
  if (cached_file == sourcefile && cached_array && cached_array->len > sourceline) {
    cached_array->ptr[sourceline]++;
    return;
  }

  cached_file = sourcefile;
  cached_array = coverage_increase_counter_uncached(sourcefile, sourceline, 0);
}

static void coverage_event_coverage_hook(rb_event_t event, NODE *node, VALUE self, ID mid, VALUE klass) {
  char *sourcefile;
  unsigned int sourceline;
  static unsigned int in_hook = 0;

  if (in_hook) return;

  in_hook++;

#if COVERAGE_DEBUG_EVENTS
  do {
    int status;
    VALUE old_exception;
    old_exception = rb_gv_get("$!");
    rb_protect(rb_inspect, klass, &status);

    if (!status) {
      printf("EVENT: %d %s %s %s %d\n", event,
             klass ? RSTRING(rb_inspect(klass))->ptr : "",
             mid ? (mid == ID_ALLOCATOR ? "ID_ALLOCATOR" : rb_id2name(mid))
             : "unknown",
             node ? node->nd_file : "", node ? nd_line(node) : 0);
    } else {
      printf("EVENT: %d %s %s %d\n", event,
             mid ? (mid == ID_ALLOCATOR ? "ID_ALLOCATOR" : rb_id2name(mid))
             : "unknown",
             node ? node->nd_file : "", node ? nd_line(node) : 0);
    }

    rb_gv_set("$!", old_exception);
  } while (0);
#endif

  if (event & RUBY_EVENT_C_CALL) coverage_mark_caller();

  if (event & (RUBY_EVENT_C_CALL | RUBY_EVENT_C_RETURN | RUBY_EVENT_CLASS)) {
    in_hook--;
    return;
  }

  if (node == NULL) {
    in_hook--;
    return;
  }

  sourcefile = node->nd_file;
  sourceline = nd_line(node) - 1;
  coverage_increase_counter_cached(sourcefile, sourceline);

  if (event & RUBY_EVENT_CALL) coverage_mark_caller();
  in_hook--;
}

static VALUE cov_install_coverage_hook(VALUE self) {
  if (!coverage_hook_set_p) {
    if (!coverinfo) coverinfo = st_init_strtable();
    coverage_hook_set_p = 1;
    /* TODO: allow C_CALL too, since it's supported already
     * the overhead is around ~30%, tested on typo */
    rb_add_event_hook(coverage_event_coverage_hook, RUBY_EVENT_ALL & ~RUBY_EVENT_C_CALL & ~RUBY_EVENT_C_RETURN & ~RUBY_EVENT_CLASS);

    return Qtrue;
  } else return Qfalse;
}

static int populate_cover(st_data_t key, st_data_t value, st_data_t cover) {
  VALUE rcover;
  VALUE rkey;
  VALUE rval;
  struct cov_array *carray;
  unsigned int i;

  rcover = (VALUE)cover;
  carray = (struct cov_array *) value;
  rkey = rb_str_new2((char*) key);
  rval = rb_ary_new2(carray->len);

  for (i = 0; i < carray->len; i++) {
    RARRAY(rval)->ptr[i] = UINT2NUM(carray->ptr[i]);
  }

  RARRAY(rval)->len = carray->len;
  rb_hash_aset(rcover, rkey, rval);

  return ST_CONTINUE;
}

static int free_table(st_data_t key, st_data_t value, st_data_t ignored) {
  struct cov_array *carray;

  carray = (struct cov_array *) value;
  free((char *)key);
  free(carray->ptr);
  free(carray);

  return ST_CONTINUE;
}

static VALUE cov_remove_coverage_hook(VALUE self) {
  if (!coverage_hook_set_p) {
    return Qfalse;
  } else {
    rb_remove_event_hook(coverage_event_coverage_hook);
    coverage_hook_set_p = 0;
    return Qtrue;
  }
}

static VALUE cov_generate_coverage_info(VALUE self) {
  VALUE cover;

  if (rb_const_defined_at(mRCOV__, id_cover)) rb_mod_remove_const(mRCOV__, ID2SYM(id_cover));

  cover = rb_hash_new();

  if (coverinfo) st_foreach(coverinfo, populate_cover, cover);

  rb_define_const(mRCOV__, "COVER", cover);

  return cover;
}

static VALUE cov_reset_coverage(VALUE self) {
  if (coverage_hook_set_p) {
    rb_raise(rb_eRuntimeError, "Cannot reset the coverage info in the middle of a traced run.");
    return Qnil;
  }

  cached_array = 0;
  cached_file = 0;
  st_foreach(coverinfo, free_table, Qnil);
  st_free_table(coverinfo);
  coverinfo = 0;

  return Qnil;
}

static VALUE cov_ABI(VALUE self) {
  VALUE ret;

  ret = rb_ary_new();
  rb_ary_push(ret, INT2FIX(RCOVRT_VERSION_MAJOR));
  rb_ary_push(ret, INT2FIX(RCOVRT_VERSION_MINOR));
  rb_ary_push(ret, INT2FIX(RCOVRT_VERSION_REV));

  return ret;
}

void Init_rcovrt() {
  ID id_rcov = rb_intern("Rcov");
  ID id_coverage__ = rb_intern("RCOV__");
  ID id_script_lines__ = rb_intern("SCRIPT_LINES__");

  id_cover = rb_intern("COVER");

  if (rb_const_defined(rb_cObject, id_rcov)) {
    mRcov = rb_const_get(rb_cObject, id_rcov);
  } else {
    mRcov = rb_define_module("Rcov");
  }

  if (rb_const_defined(mRcov, id_coverage__)) {
    mRCOV__ = rb_const_get_at(mRcov, id_coverage__);
  } else {
    mRCOV__ = rb_define_module_under(mRcov, "RCOV__");
  }

  if (rb_const_defined(rb_cObject, id_script_lines__)) {
    oSCRIPT_LINES__ = rb_const_get(rb_cObject, rb_intern("SCRIPT_LINES__"));
  } else {
    oSCRIPT_LINES__ = rb_hash_new();
    rb_const_set(rb_cObject, id_script_lines__, oSCRIPT_LINES__);
  }

  coverage_hook_set_p = 0;

  rb_define_singleton_method(mRCOV__, "install_coverage_hook", cov_install_coverage_hook, 0);
  rb_define_singleton_method(mRCOV__, "remove_coverage_hook",  cov_remove_coverage_hook, 0);
  rb_define_singleton_method(mRCOV__, "generate_coverage_info", cov_generate_coverage_info, 0);
  rb_define_singleton_method(mRCOV__, "reset_coverage", cov_reset_coverage, 0);
  rb_define_singleton_method(mRCOV__, "ABI", cov_ABI, 0);

  Init_rcov_callsite();
}
