// Half of this file comes from contributions from Nitay Joffe (nitay@powerset.com)
// Much of the rest (almost) directly ported (or pulled) from thrift-py's fastbinary.c
// Everything else via Kevin Clark (kevin@powerset.com)
#include <stdint.h>
#include <stdbool.h>
 
#include <ruby.h>
#include <st.h>
#include <netinet/in.h>
 
// #define __DEBUG__
 
#ifndef HAVE_STRLCPY
 
static
size_t
strlcpy (char *dst, const char *src, size_t dst_sz)
{
    size_t n;
 
    for (n = 0; n < dst_sz; n++) {
      if ((*dst++ = *src++) == '\0')
        break;
    }
 
    if (n < dst_sz)
      return n;
    if (n > 0)
      *(dst - 1) = '\0';
    return n + strlen (src);
}
 
#endif
 
#define dbg() fprintf(stderr, "%s:%d\n", __FUNCTION__, __LINE__)
 
 
// TODO (kevinclark): This was here from the patch/python. Not sure
// If it's actually that big a pain. Need to look into pulling
// From the right place
 
// Stolen out of TProtocol.h.
// It would be a huge pain to have both get this from one place.
 
enum TType {
  T_STOP       = 0,
  T_BOOL       = 2,
  T_BYTE       = 3,
  T_I16        = 6,
  T_I32        = 8,
  T_I64        = 10,
  T_DBL        = 4,
  T_STR        = 11,
  T_STRCT      = 12,
  T_MAP        = 13,
  T_SET        = 14,
  T_LIST       = 15
  // T_VOID       = 1,
  // T_I08        = 3,
  // T_U64        = 9,
  // T_UTF7       = 11,
  // T_UTF8       = 16,
  // T_UTF16      = 17
};
 
#define IS_CONTAINER(x) (x == T_MAP || x == T_SET || x == T_LIST)
 
// Same comment as the enum.  Sorry.
#ifdef HAVE_ENDIAN_H
#include <endian.h>
#endif
 
#ifndef __BYTE_ORDER
# if defined(BYTE_ORDER) && defined(LITTLE_ENDIAN) && defined(BIG_ENDIAN)
#  define __BYTE_ORDER BYTE_ORDER
#  define __LITTLE_ENDIAN LITTLE_ENDIAN
#  define __BIG_ENDIAN BIG_ENDIAN
# else
#  error "Cannot determine endianness"
# endif
#endif
 
#if __BYTE_ORDER == __BIG_ENDIAN
# define ntohll(n) (n)
# define htonll(n) (n)
#elif __BYTE_ORDER == __LITTLE_ENDIAN
# if defined(__GNUC__) && defined(__GLIBC__)
#  include <byteswap.h>
#  define ntohll(n) bswap_64(n)
#  define htonll(n) bswap_64(n)
# else /* GNUC & GLIBC */
#  define ntohll(n) ( (((unsigned long long)ntohl(n)) << 32) + ntohl(n >> 32) )
#  define htonll(n) ( (((unsigned long long)htonl(n)) << 32) + htonl(n >> 32) )
# endif /* GNUC & GLIBC */
#else /* __BYTE_ORDER */
# error "Can't define htonll or ntohll!"
#endif
 
 
// -----------------------------------------------------------------------------
// Cached interned strings and such
// -----------------------------------------------------------------------------
 
static VALUE class_tbpa;
static ID type_sym;
static ID class_sym;
static ID key_sym;
static ID value_sym;
static ID element_sym;
static ID name_sym;
static ID fields_id;
static ID consume_bang_id;
static ID string_buffer_id;
static ID borrow_id;
static ID keys_id;
 
static const uint32_t VERSION_MASK = 0xffff0000;
static const uint32_t VERSION_1 = 0x80010000;
 
// -----------------------------------------------------------------------------
// Structs so I don't have to keep calling rb_hash_aref
// -----------------------------------------------------------------------------
 
// { :type => field[:type],
//   :class => field[:class],
//   :key => field[:key],
//   :value => field[:value],
//   :element => field[:element] }
 
struct _thrift_map;
struct _field_spec;
 
typedef union {
  VALUE class;
  struct _thrift_map* map;
  struct _field_spec* element;
} container_data;
 
typedef struct _field_spec {
  int type;
  char* name;
  container_data data;
} field_spec;
 
typedef struct _thrift_map {
  field_spec* key;
  field_spec* value; 
} thrift_map;
 
 
static void free_field_spec(field_spec* spec) {
  switch(spec->type) {
    case T_LIST:
    case T_SET:
      free_field_spec(spec->data.element);
      break;
    
    case T_MAP:
      free_field_spec(spec->data.map->key);
      free_field_spec(spec->data.map->value);
      free(spec->data.map);
      break;
  }
  
  free(spec);
}
 
// Parses a ruby field spec into a C struct
//
// Simple fields look like:
// { :name => .., :type => .. }
// Structs add the :class attribute
// Maps adds :key and :value attributes, field specs
// Lists and Sets add an :element, a field spec
static field_spec* parse_field_spec(VALUE field_data) {
  int type = NUM2INT(rb_hash_aref(field_data, type_sym));
  VALUE name = rb_hash_aref(field_data, name_sym);
  field_spec* spec = (field_spec *) malloc(sizeof(field_spec));
 
#ifdef __DEBUG__ // No need for this in prod since I set all the fields
  bzero(spec, sizeof(field_spec));
#endif
 
  spec->type = type;
  
  if (Qnil != name) {
    spec->name = StringValuePtr(name);
  } else {
    spec->name = NULL;
  }
  
  switch(type) {
    case T_STRCT: {
      spec->data.class = rb_hash_aref(field_data, class_sym);
      break;
    }
    
    case T_MAP: {
      VALUE key_fields = rb_hash_aref(field_data, key_sym);
      VALUE value_fields = rb_hash_aref(field_data, value_sym);
      thrift_map* map = (thrift_map *) malloc(sizeof(thrift_map));
      
      map->key = parse_field_spec(key_fields);
      map->value = parse_field_spec(value_fields);
      spec->data.map = map;
      
      break;
    }
    
    case T_LIST: 
    case T_SET:
    {
      VALUE list_fields = rb_hash_aref(field_data, element_sym);
      spec->data.element = parse_field_spec(list_fields);
      break;
    }
  }
  
  return spec;
}
 
 
// -----------------------------------------------------------------------------
// Serialization routines
// -----------------------------------------------------------------------------
 
 
// write_*(VALUE buf, ...) takes a value and adds it to a Ruby string buffer,
// in network order
static void write_byte(VALUE buf, int8_t val) {
  rb_str_buf_cat(buf, (char*)&val, sizeof(int8_t));
}
 
static void write_i16(VALUE buf, int16_t val) {
  int16_t net = (int16_t)htons(val);
  rb_str_buf_cat(buf, (char*)&net, sizeof(int16_t));
}
 
static void write_i32(VALUE buf, int32_t val) {
  int32_t net = (int32_t)htonl(val);
  rb_str_buf_cat(buf, (char*)&net, sizeof(int32_t));
}
 
static void write_i64(VALUE buf, int64_t val) {
  int64_t net = (int64_t)htonll(val);
  rb_str_buf_cat(buf, (char*)&net, sizeof(int64_t));
}
 
static void write_double(VALUE buf, double dub) {
  // Unfortunately, bitwise_cast doesn't work in C.  Bad C!
  union {
    double f;
    int64_t t;
  } transfer;
  transfer.f = dub;
  write_i64(buf, transfer.t);
}
 
static void write_string(VALUE buf, char* str) {
  int32_t len = strlen(str);
  write_i32(buf, len);
  rb_str_buf_cat2(buf, str);
}
 
// Some functions macro'd out because they're nops for the binary protocol
// but placeholders were desired in case things change
#define write_struct_begin(buf)
#define write_struct_end(buf)
 
static void write_field_begin(VALUE buf, char* name, int type, int fid) {
#ifdef __DEBUG__
  fprintf(stderr, "Writing field beginning: %s %d %d\n", name, type, fid);
#endif
 
  write_byte(buf, (int8_t)type);
  write_i16(buf, (int16_t)fid);
}
 
#define write_field_end(buf)
 
static void write_field_stop(VALUE buf) {
  write_byte(buf, T_STOP);
}
 
static void write_map_begin(VALUE buf, int8_t ktype, int8_t vtype, int32_t sz) {
  write_byte(buf, ktype);
  write_byte(buf, vtype);
  write_i32(buf, sz);
}
 
#define write_map_end(buf);
 
static void write_list_begin(VALUE buf, int type, int sz) {
  write_byte(buf, type);
  write_i32(buf, sz);
}
 
#define write_list_end(buf)
 
static void write_set_begin(VALUE buf, int type, int sz) {
  write_byte(buf, type);
  write_i32(buf, sz);
}
 
#define write_set_end(buf)
 
static void binary_encoding(VALUE buf, VALUE obj, int type);
 
// Handles container types: Map, Set, List
static void write_container(VALUE buf, VALUE value, field_spec* spec) {
  int sz, i;
  
  switch(spec->type) {
    case T_MAP: {
      VALUE keys;
      VALUE key;
      VALUE val;
      
      keys = rb_funcall(value, keys_id, 0);
      
      sz = RARRAY(keys)->len;
      
      write_map_begin(buf, spec->data.map->key->type, spec->data.map->value->type, sz);
      
      for (i = 0; i < sz; i++) {
        key = rb_ary_entry(keys, i);
        val = rb_hash_aref(value, key);
        
        if (IS_CONTAINER(spec->data.map->key->type)) {
          write_container(buf, key, spec->data.map->key);
        } else {
          binary_encoding(buf, key, spec->data.map->key->type);
        }
        
        if (IS_CONTAINER(spec->data.map->value->type)) {
          write_container(buf, val, spec->data.map->value);
        } else {
          binary_encoding(buf, val, spec->data.map->value->type);
        }
      }
      
      write_map_end(buf);
 
      break;
    }
    
    case T_LIST: {
      sz = RARRAY(value)->len;
      
      write_list_begin(buf, spec->data.element->type, sz);
      for (i = 0; i < sz; ++i) {
        if (IS_CONTAINER(spec->data.element->type)) {
          write_container(buf, rb_ary_entry(value, i), spec->data.element);
        } else {
          binary_encoding(buf, rb_ary_entry(value, i), spec->data.element->type);
        }
      }
      write_list_end(buf);
      break;
    }
 
    case T_SET: {
      VALUE items;
      
      if (TYPE(value) == T_ARRAY) {
        items = value;
      } else {
        Check_Type(value, T_HASH);
        items = rb_funcall(value, keys_id, 0);
      }
 
      sz = RARRAY(items)->len;
      
      write_set_begin(buf, spec->data.element->type, sz);
      
      for (i = 0; i < sz; i++) {
        if (IS_CONTAINER(spec->data.element->type)) {
          write_container(buf, rb_ary_entry(items, i), spec->data.element);
        } else {
          binary_encoding(buf, rb_ary_entry(items, i), spec->data.element->type);
        }
      }
      
      write_set_end(buf);
      break;
    }
  }
}
 
// Takes the field id, data to be encoded, buffer and enclosing object
// to be encoded. buf and obj passed as a ruby array for rb_hash_foreach.
// TODO(kevinclark): See if they can be passed individually to avoid object
// creation
static int encode_field(VALUE fid, VALUE data, VALUE ary) {
  field_spec *spec = parse_field_spec(data);
  
  VALUE buf = rb_ary_entry(ary, 0);
  VALUE obj = rb_ary_entry(ary, 1);
  char name_buf[128];
  
  name_buf[0] = '@';
  strlcpy(&name_buf[1], spec->name, sizeof(name_buf) - 1);
  
  VALUE value = rb_ivar_get(obj, rb_intern(name_buf));
  
  if (Qnil == value) {
    free_field_spec(spec);
    return 0;
  }
     
  write_field_begin(buf, spec->name, spec->type, NUM2INT(fid));
  
  if (IS_CONTAINER(spec->type)) {
    write_container(buf, value, spec);
  } else {
    binary_encoding(buf, value, spec->type);
  }
  write_field_end(buf);
  
  free_field_spec(spec);
  
  return 0;
}
 
// -----------------------------------------------------------------------------
// TFastBinaryProtocol's main encoding loop
// -----------------------------------------------------------------------------
 
static void binary_encoding(VALUE buf, VALUE obj, int type) {
#ifdef __DEBUG__
  rb_p(rb_str_new2("Encoding binary (buf, obj, type)"));
  rb_p(rb_inspect(buf));
  rb_p(rb_inspect(obj));
  rb_p(rb_inspect(INT2FIX(type)));
#endif
 
  switch(type) {
    case T_BOOL:
      if RTEST(obj) {
        write_byte(buf, 1);
      }
      else {
        write_byte(buf, 0);
      }
      
      break;
    
    case T_BYTE:
      write_byte(buf, NUM2INT(obj));
      break;
    
    case T_I16:
      write_i16(buf, NUM2INT(obj));
      break;
    
    case T_I32:
      write_i32(buf, NUM2INT(obj));
      break;
    
    case T_I64:
      write_i64(buf, rb_num2ll(obj));
      break;
    
    case T_DBL:
      write_double(buf, NUM2DBL(obj));
      break;
 
    case T_STR:
      write_string(buf, StringValuePtr(obj));
      break;
          
    case T_STRCT: {
      // rb_hash_foreach has to take args as a ruby array
      VALUE args = rb_ary_new3(2, buf, obj);