#include <lua.h>
#include <lauxlib.h>
#include "moonhash_sha3.c"

void moonhash_push_hex(lua_State *L, unsigned char *buf, int len) {
  int i;
  unsigned char n;
  char str[2*len+1];
  for (i=0; i<len; i++) {
    n = buf[i] >> 4;
    str[2*i+0] = n + ((n < 10) ? '0' : ('a' - 10));
    n = buf[i] & 0x0f;
    str[2*i+1] = n + ((n < 10) ? '0' : ('a' - 10));
  }
  str[2*len] = 0;
  lua_pushstring(L, str);
}

typedef void (*moonhash_sha3_fptr)(const uint8_t *, uint64_t, uint8_t *);
typedef void (*moonhash_shake_fptr)(const uint8_t *, uint64_t, uint8_t *, uint64_t);

int moonhash_sha3(lua_State *L, moonhash_sha3_fptr hashfunc, int len) {
  const char *input;
  size_t inputlen;
  unsigned char output[len];
  input = luaL_checklstring(L, 1, &inputlen);
  hashfunc((const uint8_t *)input, inputlen, output);
  moonhash_push_hex(L, output, len);
  return 1;
}

int moonhash_sha3_224(lua_State *L) {
  return moonhash_sha3(L, FIPS202_SHA3_224, 224/8);
}
int moonhash_sha3_256(lua_State *L) {
  return moonhash_sha3(L, FIPS202_SHA3_256, 256/8);
}
int moonhash_sha3_384(lua_State *L) {
  return moonhash_sha3(L, FIPS202_SHA3_384, 384/8);
}
int moonhash_sha3_512(lua_State *L) {
  return moonhash_sha3(L, FIPS202_SHA3_512, 512/8);
}

int moonhash_shake(lua_State *L, int R, lua_Integer deflen) {
  const char *input;
  size_t inputlen;
  lua_Integer outputlen;
  const char *alphabet;
  size_t alen;
  int abits = 0;
  uint8_t s[200];
  luaL_Buffer luabuf;
  uint8_t *output;
  int readpos = 0;
  lua_Integer writepos = 0;
  int rbits = 0;
  int rbuf = 0;
  int rvalue;
  input = luaL_checklstring(L, 1, &inputlen);
  outputlen = luaL_optinteger(L, 2, deflen);
  alphabet = luaL_optlstring(L, 3, "0123456789abcdef", &alen);
  luaL_argcheck(L, alen>1, 3, "too few characters in alphabet");
  {
    size_t t = alen-1;
    while (t) {
      abits++;
      if (abits > 8) luaL_argcheck(L, 0, 3, "too many characters in alphabet");
      t >>= 1;
    }
  }
  KeccakF1600Init(s, R, (const uint8_t *)input, inputlen, 0x1F);
  output = (uint8_t *)luaL_buffinitsize(L, &luabuf, outputlen);
  while (writepos < outputlen) {
    if (rbits < abits) {
      if (readpos == R) {
        KeccakF1600(s);
        readpos = 0;
      }
      rbuf = (rbuf << 8) | s[readpos++];
      rbits += 8;
    }
    rbits -= abits;
    rvalue = rbuf >> rbits;
    rbuf &= (1<<rbits)-1;
    if (rvalue < alen) output[writepos++] = alphabet[rvalue];
  }
  luaL_pushresultsize(&luabuf, outputlen);
  return 1;
}

int moonhash_shake128(lua_State *L) {
  return moonhash_shake(L, (1600-2*128)/8, 128/4);
}

int moonhash_shake256(lua_State *L) {
  return moonhash_shake(L, (1600-2*256)/8, 256/4);
}

static const struct luaL_Reg moonhash_module_functions[] = {
 {"sha3_224", moonhash_sha3_224},
 {"sha3_256", moonhash_sha3_256},
 {"sha3_384", moonhash_sha3_384},
 {"sha3_512", moonhash_sha3_512},
 {"shake128", moonhash_shake128},
 {"shake256", moonhash_shake256},
 {NULL, NULL}
};

int luaopen_moonhash(lua_State *L) {
  lua_newtable(L);
  luaL_setfuncs(L, moonhash_module_functions, 0);
  return 1;
}
