#include <stdio.h>
#include <ruby.h>
#include <node.h>
#include <version.h>

/* <DebugStuff> use the --debug argument to ruby to enable */
const char *node_names[] = {"METHOD", "FBODY", "CFUNC", "SCOPE", "BLOCK", 
  "IF", "CASE", "WHEN", "OPT_N", "WHILE", "UNTIL", "ITER", "FOR", "BREAK", 
  "NEXT", "REDO", "RETRY", "BEGIN", "RESCUE", "RESBODY", "ENSURE", "AND", 
  "OR", "NOT", "MASGN", "LASGN", "DASGN", "DASGN_CURR", "GASGN", "IASGN", 
  "CDECL", "CVASGN", "CVDECL", "OP_ASGN1", "OP_ASGN2", "OP_ASGN_AND", 
  "OP_ASGN_OR", "CALL", "FCALL", "VCALL", "SUPER", "ZSUPER", "ARRAY", 
  "ZARRAY", "HASH", "RETURN", "YIELD", "LVAR", "DVAR", "GVAR", "IVAR", 
  "CONST", "CVAR", "NTH_REF", "BACK_REF", "MATCH", "MATCH2", "MATCH3", 
  "LIT", "STR", "DSTR", "XSTR", "DXSTR", "EVSTR", "DREGX", "DREGX_ONCE", 
  "ARGS", "ARGSCAT", "ARGSPUSH", "SPLAT", "TO_ARY", "SVALUE", "BLOCK_ARG", 
  "BLOCK_PASS", "DEFN", "DEFS", "ALIAS", "VALIAS", "UNDEF", "CLASS", "MODULE",
  "SCLASS", "COLON2", "COLON3", "CREF", "DOT2", "DOT3", "FLIP2", "FLIP3", 
  "ATTRSET", "SELF", "NIL", "TRUE", "FALSE", "DEFINED", "NEWLINE", "POSTEXE",
  "ALLOCA", "DMETHOD", "BMETHOD", "MEMO", "IFUNC", "DSYM", "ATTRASGN", "LAST"
};

void inspect_node(NODE *n) {
  if(n) {
    fprintf(stderr, "node %s: %s%s%s\n",
      node_names[nd_type(n)],
      (RNODE(n)->u1.node != NULL ? "u1 " : "   "),
      (RNODE(n)->u2.node != NULL ? "u2 " : "   "),
      (RNODE(n)->u3.node != NULL ? "u3 " : "   "));
  } else {
    fprintf(stderr, "null node!\n");
  }
}
/* </DebugStuff> */

NODE* get_arg_node(NODE *body) {
  NODE *node = NULL, *block, *args;
  
  if(RTEST(ruby_debug)) { inspect_node(body); } 
  if(!body) { return NULL; }  
  
  block = body->nd_next;
  if(RTEST(ruby_debug)) { inspect_node(block); }
  if(!block) { return NULL; }
  
  args = block->nd_head;
  if(RTEST(ruby_debug)) { inspect_node(args); }
  if(!args) { return NULL; }
  
  return args;
}

NODE* get_body_node(VALUE method) { 
#if RUBY_VERSION_CODE < 190
  /* Exposing struct METHOD is terrible but allows for easily accessing  body.
   * We're moving to 1.9 soon so we don't need to do this in the future */
  struct METHOD { 
    VALUE klass, rklass;
    VALUE recv;
    ID id, oid;
    int safe_level;
    NODE *body;
  };
  struct METHOD *data;
  Data_Get_Struct(method, struct METHOD, data);
  return data->body;
#else
  /* 1.9 has Method#receiver. use it instead of peeking at secret structures 
   */
  ID id;
  VALUE recv;
  recv = rb_funcall(method, rb_intern('receiver'), 0);
  id = rb_to_id(recv);
  st_lookup(RCLASS(recv)->m_tbl, id, &n);
  return (NODE*)n;
#endif
}

VALUE required_args(ID* locals, int max_args) {
  int i;
  VALUE arg_array = rb_ary_new();  
  for (i = 0; i < max_args; i++) {
    rb_ary_push(arg_array, ID2SYM(locals[i + 3]));
  }
  
  if(RTEST(ruby_debug)) {
    fprintf(stderr, "required_args: %s\n",
      RSTRING(rb_funcall(arg_array, rb_intern("inspect"), 0))->ptr);
  }
  
  return arg_array;
}

VALUE optional_args(NODE *args, ID *locals, int max_args) {
  NODE *optnode;
  int i;
  VALUE arg_array = rb_ary_new();  
  
  i =  max_args; 
  for(optnode = args->nd_opt; optnode; optnode = optnode->nd_next) {
    rb_ary_push(arg_array, ID2SYM(locals[i + 3]));
    i++;
  }
  
  if(RTEST(ruby_debug)) {
    fprintf(stderr, "optional_args: %s\n",
      RSTRING(rb_funcall(arg_array, rb_intern("inspect"), 0))->ptr);
  }
  
  return arg_array;
}

VALUE optional_args_length(NODE *args) {
  NODE *optnode;
  int l=0;
  
  for(optnode = args->nd_opt; optnode; optnode = optnode->nd_next) { l++; }
  return l;
}

/* arg_index should be optional_args.length + max_args */
VALUE splat_arg(NODE *args, ID *locals, int max_args) {
  long arg_count = (long)args->nd_rest;
  int arg_index = optional_args_length(args) + max_args;
  
  if(RTEST(ruby_debug)) {
    fprintf(stderr, "splat_arg: arg_count = %ld\n", arg_count);
    fprintf(stderr, "splat_arg: arg_index = %d\n", arg_index);
  } 
  
  /* splat argument */
  if (arg_count > 0 && locals[arg_index + 3]) {
    if(RTEST(ruby_debug)) { fprintf(stderr, "hello world\n"); }
    return ID2SYM(locals[arg_index + 3]);
  } else {
    return Qnil;
  }
}

VALUE rb_method_args_p(VALUE self) {
  return get_arg_node(get_body_node(self)) ? Qtrue : Qfalse;
}

VALUE rb_method_args(VALUE self) {
  NODE *args, *body = get_body_node(self);    
  ID *locals;
  VALUE array, splat;
  int max_args;
  
  if(rb_method_args_p(self) == Qfalse) { 
    rb_raise(rb_eStandardError, "Cannot find arguments for this method");    
  }
    
  args = get_arg_node(body);
  locals = body->nd_tbl;
  max_args = args->nd_cnt;
  
  array = required_args(locals, max_args);
  rb_ary_concat(array, optional_args(args, locals, max_args));
  if((splat = splat_arg(args, locals, max_args)) != Qnil) {
    rb_ary_push(array, splat);
  }
  
  return array;
}

VALUE rb_method_required_args(VALUE self) {
  NODE *args, *body;
  
  if(rb_method_args_p(self) == Qfalse) { 
    rb_raise(rb_eStandardError, "Cannot find arguments for this method");    
  }
    
  body = get_body_node(self);
  args = get_arg_node(body);
  
  return required_args(body->nd_tbl, args->nd_cnt);
}

VALUE rb_method_optional_args(VALUE self) {
  NODE *args, *body;
  
  if(rb_method_args_p(self) == Qfalse) { 
    rb_raise(rb_eStandardError, "Cannot find arguments for this method");    
  }  
  body = get_body_node(self);
  args = get_arg_node(body);
  
  return optional_args(get_arg_node(body), body->nd_tbl, args->nd_cnt);  
}

VALUE rb_method_splat_arg(VALUE self) {
  NODE *args, *body = get_body_node(self);    
  ID *locals;
  int max_args; 
  
  if(rb_method_args_p(self) == Qfalse) { 
    rb_raise(rb_eStandardError, "Cannot find arguments for this method");    
  }
  
  args = get_arg_node(body);
  locals = body->nd_tbl;
  max_args = args->nd_cnt;
  
  return splat_arg(args, locals, max_args);
}

void Init_method_args() {
  VALUE cMethod = rb_const_get(rb_cObject, rb_intern("Method"));
  VALUE cUBMethod = rb_const_get(rb_cObject, rb_intern("UnboundMethod"));
  
  rb_define_method(cMethod, "args?", rb_method_args_p, 0);
  rb_define_method(cMethod, "args", rb_method_args, 0);
  rb_define_method(cMethod, "required_args", rb_method_required_args, 0);
  rb_define_method(cMethod, "optional_args", rb_method_optional_args, 0);
  rb_define_method(cMethod, "splat_arg", rb_method_splat_arg, 0);
  // rb_define_method(rb_cMethod, "block_arg", rb_method_block_arg, 0);
  
  rb_define_method(cUBMethod, "args?", rb_method_args_p, 0);
  rb_define_method(cUBMethod, "args", rb_method_args, 0);
  rb_define_method(cUBMethod, "required_args", rb_method_required_args, 0);
  rb_define_method(cUBMethod, "optional_args", rb_method_optional_args, 0);
  rb_define_method(cUBMethod, "splat_arg", rb_method_splat_arg, 0);
}