#include <ctype.h>
#include <string.h>
#include <stdint.h>

#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

/*
Description of the 'format' field.
- (...) are comments.
- whitespace is ignored.
- numbers are numeric arguments, noted as "*" below. unless otherwise stated,
  these are the number of repetitions of the next command. "4 i" == "iiii"
- native types:
  "f" float
  "F" double
  "n" lua_Integer
  "N" lua_Number
  "h"/"H" char     (uppercase = unsigned)
  "w"/"W" short
  "i"/"I" int
  "l"/"L" long
  "a"/"A" int1_t
  "b"/"B" int2_t
  "c"/"C" int4_t
  "d"/"D" int8_t
- other data: (not aligned)
  "x" padding (1-byte; ignore when unpacking, fill with zeroes when packing)
  "*s" string of given length (when packing, arg is optional)
  "z"/"0z" zero-terminated string
  "*z" zero-terminated string of given length, but with padded zeroes
    (arg is the exact number of bytes that will be read/written)
    (when packing, if arg < length, write only arg bytes)
  "St" (where t is some native integer type, e.g. "SC")
    the integer, and a string of that length
- special flags:
  ">" big endian
  "<" little endian
  "=" native endian
  "*!" alignment (when without n, use native) (currently can't be turned off)
    (n is the maximum alignment. for example, 4! is Linux-like alignment where
    even 8-byte doubles are 4-byte aligned)

-----------------

API:
 string.pack([outptr, [stopptr|maxlen]], fmt, {data}|data...) -> outptr ? none : outstr
 string.unpack(instr|inptr, [stopptr|maxlen], fmt, [bool rettable], [offset from 0]) ->
     rettable ? {data},newpos : data...,newpos

the rest was moved to SDL itself:
 rw:read(num) (num=nil means read all) -> outstr
 rw:write(str, [num]) (num=nil means whole str) -> true (or nil,error,nwrote)
 rw:seek(["end"|"set"|"cur"], [offset]) -> none
 SDL.RWFromNewMem(len) -> rw
 SDL.RWFromMem(ptr, len) -> rw
 rw:baseptr() -> ptr
 rw:hereptr() -> ptr
 rw:disown()
 SDL.IMG_Load_RW(rw, freesrc)
 SDL.LoadWAV_RW(rw, freesrc)
*/


// how many bytes is 'double' is aligned to
struct dummyalign { char c; double d; };
#define MAXALIGN (sizeof(struct dummyalign) - sizeof(double))

// FIXME: in the original struct.c from ~roberto, this is used instead:
// #define NEW_MAXALIGN  (MAXALIGN > sizeof(int) ? MAXALIGN : sizeof(int))
// why? is the current approach invalid? IMHO, if MAXALIGN < sizeof(int),
// the platform probably doesn't align at all, and MAXALIGN is correct.

struct PackState {
  int swap;
  int align;       // max alignment for the packed data
  int usestring;
  char token;
  int num_arg;
  
  const char *p_format;
  
  lua_State *L;
};

static void doswap(char *a, size_t n) {
  int i, j;
  for (i=0, j=n-1, n=n/2; n--; i++, j--) {
    char t = a[i]; a[i] = a[j]; a[j] = t;
  }
}

static int doalign(int offset, int align, int maxalign) {
  if(align > maxalign) align = maxalign;
  return (align - (offset % align)) % align;
}

static void get_next_token(struct PackState *st) {
  char input;
  st->usestring = 0;
  if(st->token == 's' || st->token == 'z') st->token = 0, st->num_arg = 0;
  if(st->num_arg && st->token) {
    st->num_arg--;
    return;
  }
  while(isspace(*st->p_format) || *st->p_format == '(') {
    if(*st->p_format == '(') {
      while(*st->p_format != ')' && *st->p_format != '\0') st->p_format++;
      if(*st->p_format == '\0') luaL_error(st->L, "comment not terminated");
    }
    st->p_format++;
  }
  if(*st->p_format == '\0') {
    if(st->num_arg) luaL_error(st->L, "trailing number at end of format");
    st->token = -1; return;   // end of format string
  }
  if(*st->p_format >= '0' && *st->p_format <= '9') {
    st->token = 0;
    st->num_arg = *st->p_format - '0';
    st->p_format++;
    while(*st->p_format >= '0' && *st->p_format <= '9') {
      st->num_arg = 10*st->num_arg + *st->p_format-'0';
      st->p_format++;
    }
    get_next_token(st); return;
  }
  input = *st->p_format;
  st->p_format++;
  if(strchr("S=<>", input) && st->num_arg != 0)
    luaL_error(st->L, "unexpected number before '%c'", input);
  if(input == '!') {   // used alignment
    st->align = (st->num_arg ? st->num_arg : MAXALIGN);
    st->num_arg = 0;
    get_next_token(st); return;
  }
  if(input == '=') {   // native endian, don't swap
    st->swap = 0;
    get_next_token(st); return;
  }
  if(input == '<') {   // little-endian packed
    int arch_is_le = 1;
    st->swap = !(*(char *)&arch_is_le);
    get_next_token(st); return;
  }
  if(input == '>') {   // big-endian packed
    int arch_is_le = 1;
    st->swap = (*(char *)&arch_is_le);
    get_next_token(st); return;
  }
  if(input == 'S') {   // string with prepended packed length
    get_next_token(st);
    if(!strchr("hHwWiIlLaAbBcCdD", st->token))
      luaL_error(st->L, "expected integer type after 'S'");
    if(st->num_arg)
      luaL_error(st->L, "unexpected number after 'S'");
    st->usestring = 1;
    return;
  }
  if(strchr("zxsfFnNhHwWiIlLaAbBcCdD", input)) {   // supported chars
    if(st->num_arg != 0 && input != 's' && input != 'z')
      st->num_arg--;
    st->token = input;
    return;
  }
  luaL_error(st->L, "unexpected character '%c'", input);
}

#define UNPACK(OP,T)                                                 \
  case OP: {                                                         \
    T val; int size = sizeof(val);                                   \
    if(st.align && !st.usestring)                                    \
      pos += doalign(pos, size, st.align);                           \
    if(maxlen != -1 && pos+size > maxlen)                            \
      luaL_error(L, "unexpected end of input at '%c'", st.token);    \
    memcpy(&val, data+pos, size);                                    \
    pos += size;                                                     \
    if(st.swap) doswap((char *)&val, size);                          \
    if(st.usestring) {                                               \
      if(maxlen != -1 && pos+val > maxlen)                           \
        luaL_error(L, "unexpected end of input in 'S' string");      \
      lua_pushlstring(L, data+pos, val);                             \
      pos += val;                                                    \
    }                                                                \
    else                                                             \
      lua_pushnumber(L, (lua_Number)val);                            \
    break;                                                           \
  }

static int str_unpack(lua_State *L) {
  const char *data = NULL;
  int offset = 0, pos = 0, fmtarg = 2, intable = 0;
  size_t maxlen = (size_t)-1;
  struct PackState st;
  memset(&st, 0, sizeof st);
  st.L = L;
  if(lua_type(L, 1) == LUA_TLIGHTUSERDATA)
    data = lua_touserdata(L, 1);
  else
    data = luaL_checklstring(L, 1, &maxlen);
  if(lua_type(L, 1) == LUA_TLIGHTUSERDATA &&
     lua_type(L, 2) == LUA_TLIGHTUSERDATA) {
    const char *stop = lua_touserdata(L, 2);
    if(stop < data) luaL_error(L, "stop pointer is sooner than start pointer");
    maxlen = stop - data; fmtarg++;
  }
  else if(lua_isnumber(L, 2)) {
    maxlen = luaL_checkinteger(L, 2); fmtarg++;
  }
  st.p_format = luaL_checkstring(L, fmtarg);
  if(lua_isboolean(L, fmtarg+1) && lua_toboolean(L, fmtarg+1)) {
    intable = 1; fmtarg++;
  }
  if(lua_isnumber(L, fmtarg+1)) {
    offset = luaL_checkinteger(L, fmtarg+1);
    if(maxlen != -1 && offset > maxlen)
      luaL_error(L, "unexpected end of input at offset");
    // offset isn't just an increment to pos because alignment is computed
    // against pos. this way, pos is always 0 at the beginning.
    data += offset;
    if(maxlen != -1) maxlen -= offset;
  }
  lua_settop(L, 0);
  if(intable)
    lua_newtable(L);
  while(1) {
    get_next_token(&st);
    if(st.token == -1) return lua_pushinteger(L, pos+offset), lua_gettop(L);
    if(st.token != 'x')
      luaL_checkstack(L, 1, "too many results");
    switch(st.token) {
      case 'x':
        if(maxlen != -1 && pos+1 > maxlen)
          luaL_error(L, "unexpected end of input at '%c'", st.token);
        pos++;
        break;
        
      case 'z':
        if(st.num_arg == 0) {
          while((maxlen == -1 || pos+st.num_arg < maxlen) &&
              data[pos+st.num_arg] != 0)
            st.num_arg++;
          st.num_arg++; // the "\0" itself
        }
        // no break
      case 's':
        if(maxlen != -1 && pos+st.num_arg > maxlen)
          luaL_error(L, "unexpected end of input at '%c'", st.token);
        if(st.token == 'z')
          lua_pushstring(L, data+pos);
        else
          lua_pushlstring(L, data+pos, st.num_arg);
        pos += st.num_arg;
        break;
        
      UNPACK('f', float)
      UNPACK('F', double)
      UNPACK('n', lua_Integer)
      UNPACK('N', lua_Number)
      UNPACK('h', char)
      UNPACK('H', unsigned char)
      UNPACK('w', short)
      UNPACK('W', unsigned short)
      UNPACK('i', int)
      UNPACK('I', unsigned int)
      UNPACK('l', long)
      UNPACK('L', unsigned long)
      UNPACK('a', int8_t)
      UNPACK('A', uint8_t)
      UNPACK('b', int16_t)
      UNPACK('B', uint16_t)
      UNPACK('c', int32_t)
      UNPACK('C', uint32_t)
      UNPACK('d', int64_t)
      UNPACK('D', uint64_t)
    }
    if(st.token != 'x' && intable) {
      lua_rawseti(L, 1, intable);
      intable++;
    }
  }
}

#define OUTPUT(str, size) do {                                       \
  if(out) {                                                          \
    if(maxlen != -1 && pos+size > maxlen)                            \
      luaL_error(L, "not enough space in output at '%c'", st.token); \
    memcpy(out+pos, str, size);                                      \
  }                                                                  \
  else                                                               \
    luaL_addlstring(&b, (const char *)str, size);                    \
  pos += size;                                                       \
} while(0)

#define PACK(OP,T)                                                   \
  case OP: {                                                         \
    T val; int size = sizeof(val); const char *usedstr = NULL;       \
    if(intable) lua_rawgeti(L, intable, arg);                        \
    if(st.usestring) {                                               \
      usedstr = luaL_checklstring(L, intable ? -1 : arg, &len);      \
      val = len;                                                     \
    }                                                                \
    else val = (T)luaL_checknumber(L, intable ? -1 : arg);           \
    if(intable) lua_pop(L, 1);                                       \
    if(st.align && !st.usestring) {                                  \
      int myalign = doalign(pos, size, st.align);                    \
      while(myalign--) OUTPUT("\0", 1);                              \
    }                                                                \
    if(st.swap) doswap((char *)&val, size);                          \
    OUTPUT(&val, size);                                              \
    if(st.usestring) OUTPUT(usedstr, len);                           \
    break;                                                           \
  }

static int str_pack(lua_State *L) {
  luaL_Buffer b;
  char *out = NULL; int pos = 0, maxlen = -1;
  int intable = 0, arg = 1;
  size_t len;
  struct PackState st;
  memset(&st, 0, sizeof st);
  st.L = L;
  if(lua_type(L, arg) == LUA_TLIGHTUSERDATA) {   // outptr
    out = lua_touserdata(L, arg); arg++;
  }
  if(lua_type(L, arg) == LUA_TLIGHTUSERDATA) {   // stopptr
    maxlen = (char *)lua_touserdata(L, arg) - out; arg++;
  }
  else if(lua_type(L, arg) == LUA_TNUMBER) {
    maxlen = luaL_checkinteger(L, arg); arg++;
  }
  if(!out) {
    lua_pushnil(L);   // mark to separate arguments from the string buffer
    luaL_buffinit(L, &b);
  }
  st.p_format = luaL_checkstring(L, arg);
  if(lua_type(L, arg+1) == LUA_TTABLE) {
    intable = arg+1;
    arg = 0;
  }
  while(1) {
    get_next_token(&st);
    if(st.token == -1) break;
    if(st.token != 'x') arg++;
    switch(st.token) {
      case 'x':
        OUTPUT("\0", 1);
        break;
        
      case 's': case 'z': {
        const char *tmp;
        if(intable) lua_rawgeti(L, intable, arg);
        tmp = luaL_checklstring(L, intable ? -1 : arg, &len);
        if(intable) lua_pop(L, 1);
        if(st.num_arg != 0 && st.token == 's' && st.num_arg != len)
          luaL_error(L, "incorrect string length for 's'");
        if(st.num_arg != 0 && st.num_arg < len)
          len = st.num_arg;
        OUTPUT(tmp, len);
        if(st.token == 'z' && st.num_arg == 0)
          st.num_arg = len + 1;   // add one "\0"
        while(st.num_arg > len) {
          OUTPUT("\0", 1);
          len++;
        }
        break;
      }
      
      PACK('f', float)
      PACK('F', double)
      PACK('n', lua_Integer)
      PACK('N', lua_Number)
      PACK('h', char)
      PACK('H', unsigned char)
      PACK('w', short)
      PACK('W', unsigned short)
      PACK('i', int)
      PACK('I', unsigned int)
      PACK('l', long)
      PACK('L', unsigned long)
      PACK('a', int8_t)
      PACK('A', uint8_t)
      PACK('b', int16_t)
      PACK('B', uint16_t)
      PACK('c', int32_t)
      PACK('C', uint32_t)
      PACK('d', int64_t)
      PACK('D', uint64_t)
    }
  }
  if(out) return 0;
  luaL_pushresult(&b);
  return 1;
}

int luaopen_buf(lua_State *L) {
  const luaL_Reg reg_buf[] = {
    { "pack", str_pack },
    { "unpack", str_unpack },
    { NULL, NULL }
  };
  
  lua_settop(L, 0);
  
  luaL_register(L, LUA_STRLIBNAME, reg_buf);
  return 1;
}