/*
 * lbp.c - bp library for lua
 * API:
 *   bp.match(pat, str, [start_index]) -> nil or match_table
 *   bp.replace(pat, replacement, str, [start_index]) -> str with replacements, num_replacements
 *   for match_table in bp.matches(pat, str, [start_index]) do ... end
 *   bp.compile(pat) -> pattern object
 *       pat:match(str, [start_index])
 *       pat:replace(replacement, str, [start_index])
 *       for match in pat:matches(str, [start_index]) do ... end
 */

#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>

#include "lauxlib.h"
#include "lua.h"

#include "../match.h"
#include "../pattern.h"
#include "../printmatch.h"
#include "../utils.h"

static const char *builtins_source = (
#include "builtins.h"
);
static int MATCH_METATABLE = 0, PAT_METATABLE = 0;
static bp_pat_t *builtins;

static void push_match(lua_State *L, bp_match_t *m, const char *start);

lua_State *cur_state = NULL;

static void match_error(char **msg) {
    lua_pushstring(cur_state, *msg);
    free(*msg);
    *msg = NULL;
    lua_error(cur_state);
}

static inline void raise_parse_error(lua_State *L, maybe_pat_t m) {
    size_t err_len = (size_t)(m.value.error.end - m.value.error.start);
    char *buf = calloc(err_len + 1, sizeof(char));
    if (buf) {
        memcpy(buf, m.value.error.start, err_len);
        luaL_error(L, "%s: \"%s\"", m.value.error.msg, buf);
        free(buf);
    } else {
        luaL_error(L, "memory allocation failed");
    }
}

static int Lcompile(lua_State *L) {
    size_t patlen;
    const char *pat_text = luaL_checklstring(L, 1, &patlen);
    maybe_pat_t maybe_pat = bp_pattern(pat_text, pat_text + patlen);
    if (!maybe_pat.success) {
        raise_parse_error(L, maybe_pat);
        return 0;
    }
    bp_pat_t **pat_storage = (bp_pat_t **)lua_newuserdatauv(L, sizeof(bp_pat_t *), 1);
    *pat_storage = maybe_pat.value.pat;
    lua_pushvalue(L, 1);
    lua_setiuservalue(L, -2, 1);
    lua_pushlightuserdata(L, (void *)&PAT_METATABLE);
    lua_gettable(L, LUA_REGISTRYINDEX);
    lua_setmetatable(L, -2);
    return 1;
}

static void push_matchstring(lua_State *L, bp_match_t *m) {
    char *buf = NULL;
    size_t size = 0;
    FILE *out = open_memstream(&buf, &size);
    if (!out) luaL_error(L, "failed to allocate memory buffer");
    fprint_match(out, m->start, m, NULL);
    fflush(out);
    lua_pushlstring(L, buf, size);
    fclose(out);
}

static bp_match_t *get_first_capture(bp_match_t *m) {
    if (m->pat->type == BP_TAGGED) {
        return m;
    } else if (m->pat->type == BP_CAPTURE && !When(m->pat, BP_CAPTURE)->name) {
        return m;
    } else if (m->children) {
        for (int i = 0; m->children[i]; i++) {
            bp_match_t *cap = get_first_capture(m->children[i]);
            if (cap) return cap;
        }
    }
    return NULL;
}

static void set_capture_fields(lua_State *L, bp_match_t *m, int *n, const char *start) {
    if (m->pat->type == BP_CAPTURE) {
        bp_match_t *cap = get_first_capture(m->children[0]);
        if (!cap) cap = m->children[0];
        auto capture = When(m->pat, BP_CAPTURE);
        if (capture->namelen > 0) {
            lua_pushlstring(L, capture->name, capture->namelen);
            push_match(L, cap, start);
            lua_settable(L, -3);
        } else {
            push_match(L, cap, start);
            lua_seti(L, -2, (*n)++);
        }
    } else if (m->pat->type == BP_TAGGED) {
        push_match(L, m, start);
        lua_seti(L, -2, (*n)++);
    } else if (m->children) {
        for (int i = 0; m->children[i]; i++)
            set_capture_fields(L, m->children[i], n, start);
    }
}

static void push_match(lua_State *L, bp_match_t *m, const char *start) {
    lua_createtable(L, 1, 2);
    lua_pushlightuserdata(L, (void *)&MATCH_METATABLE);
    lua_gettable(L, LUA_REGISTRYINDEX);
    lua_setmetatable(L, -2);
    push_matchstring(L, m);
    lua_seti(L, -2, 0);

    if (m->pat->type == BP_TAGGED) {
        auto tagged = When(m->pat, BP_TAGGED);
        lua_pushlstring(L, tagged->name, tagged->namelen);
        lua_setfield(L, -2, "__tag");
    }

    int capture_num = 1;
    for (int i = 0; m->children && m->children[i]; i++)
        set_capture_fields(L, m->children[i], &capture_num, start);

    lua_pushinteger(L, 1 + (m->start - start));
    lua_setfield(L, -2, "start");
    lua_pushinteger(L, 1 + (m->end - start));
    lua_setfield(L, -2, "after");
}

static int Lmatch(lua_State *L) {
    if (lua_isstring(L, 1)) {
        if (Lcompile(L) != 1) return 0;
        lua_replace(L, 1);
    }
    bp_pat_t **at_pat = lua_touserdata(L, 1);
    bp_pat_t *pat = at_pat ? *at_pat : NULL;
    if (!pat) luaL_error(L, "Not a valid pattern");

    size_t textlen;
    const char *text = luaL_checklstring(L, 2, &textlen);
    lua_Integer index;
    if (lua_istable(L, 3)) {
        lua_getfield(L, 3, "start");
        lua_getfield(L, 3, "after");
        index = luaL_optinteger(L, -1, 1);
        if (lua_rawequal(L, -1, -2)) ++index;
    } else {
        index = luaL_optinteger(L, 3, 1);
    }
    if (index > (lua_Integer)strlen(text) + 1) return 0;

    bp_match_t *m = NULL;
    int ret = 0;
    cur_state = L;
    bp_errhand_t old = bp_set_error_handler(match_error);
    if (next_match(&m, text + index - 1, &text[textlen], pat, builtins, NULL, false)) {
        push_match(L, m, text);
        stop_matching(&m);
        ret = 1;
    }
    bp_set_error_handler(old);
    return ret;
}

static int Lreplace(lua_State *L) {
    if (lua_isstring(L, 1)) {
        if (Lcompile(L) != 1) return 0;
        lua_replace(L, 1);
    }
    bp_pat_t **at_pat = lua_touserdata(L, 1);
    bp_pat_t *pat = at_pat ? *at_pat : NULL;
    if (!pat) luaL_error(L, "Not a valid pattern");

    size_t replen, textlen;
    const char *rep_text = luaL_checklstring(L, 2, &replen);
    const char *text = luaL_checklstring(L, 3, &textlen);
    lua_Integer index = luaL_optinteger(L, 4, 1);
    if (index > (lua_Integer)strlen(text) + 1) index = (lua_Integer)strlen(text) + 1;

    maybe_pat_t maybe_replacement = bp_replacement(pat, rep_text, rep_text + replen);
    if (!maybe_replacement.success) {
        raise_parse_error(L, maybe_replacement);
        return 0;
    }

    char *buf = NULL;
    size_t size = 0;
    FILE *out = open_memstream(&buf, &size);
    int replacements = 0;
    const char *prev = text;
    bp_pat_t *rep_pat = maybe_replacement.value.pat;
    cur_state = L;
    bp_errhand_t old = bp_set_error_handler(match_error);
    for (bp_match_t *m = NULL; next_match(&m, text, &text[textlen], rep_pat, builtins, NULL, false);) {
        fwrite(prev, sizeof(char), (size_t)(m->start - prev), out);
        fprint_match(out, text, m, NULL);
        prev = m->end;
        ++replacements;
    }
    bp_set_error_handler(old);
    fwrite(prev, sizeof(char), (size_t)(&text[textlen] - prev), out);
    fflush(out);
    lua_pushlstring(L, buf, size);
    lua_pushinteger(L, replacements);
    fclose(out);

    delete_pat(&rep_pat, false);

    return 2;
}

static int iter(lua_State *L) {
    lua_geti(L, 1, 1);
    lua_geti(L, 1, 2);
    lua_replace(L, 1);
    lua_rotate(L, 1, 1);
    return Lmatch(L);
}

static int Lmatches(lua_State *L) {
    int nargs = lua_gettop(L);
    lua_pushcfunction(L, iter); // iter
    lua_createtable(L, 2, 0); // state: {pat, str}
    if (lua_isstring(L, 1)) {
        if (Lcompile(L) != 1) return 0;
    } else {
        lua_pushvalue(L, 1);
    }
    lua_seti(L, -2, 1);
    lua_pushvalue(L, 2);
    lua_seti(L, -2, 2);
    if (nargs >= 3) // start index
        lua_pushvalue(L, 3);
    else lua_pushnil(L);
    return 3;
}

static int Lmatch_tostring(lua_State *L) {
    lua_geti(L, 1, 0);
    return 1;
}

static int Lpat_source(lua_State *L) {
    lua_getiuservalue(L, 1, 1);
    return 1;
}

static int Lpat_tostring(lua_State *L) {
    luaL_Buffer b;
    luaL_buffinit(L, &b);
    luaL_addstring(&b, "Pattern [[");
    lua_getiuservalue(L, 1, 1);
    luaL_addvalue(&b);
    luaL_addstring(&b, "]]");
    luaL_pushresult(&b);
    return 1;
}

static int Lpat_gc(lua_State *L) {
    (void)L;
    bp_pat_t **at_pat = lua_touserdata(L, 1);
    bp_pat_t *pat = *at_pat;
    if (pat) delete_pat(at_pat, true);
    (void)pat;
    return 0;
}

static int Lpat_join(lua_State *L, const char *joiner) {
    if (!lua_isstring(L, 1)) {
        lua_pushcfunction(L, Lpat_source);
        lua_pushvalue(L, 1);
        lua_call(L, 1, 1);
        lua_replace(L, 1);
    }
    if (!lua_isstring(L, 2)) {
        lua_pushcfunction(L, Lpat_source);
        lua_pushvalue(L, 2);
        lua_call(L, 1, 1);
        lua_replace(L, 2);
    }

    lua_pushcfunction(L, Lcompile);
    luaL_Buffer b;
    luaL_buffinit(L, &b);
    luaL_addstring(&b, "(");
    lua_pushvalue(L, 1);
    luaL_addvalue(&b);
    luaL_addstring(&b, ")");
    luaL_addstring(&b, joiner);
    luaL_addstring(&b, "(");
    lua_pushvalue(L, 2);
    luaL_addvalue(&b);
    luaL_addstring(&b, ")");

    luaL_pushresult(&b);
    lua_call(L, 1, 1);

    return 1;
}

static int Lpat_concat(lua_State *L) { return Lpat_join(L, " "); }

static int Lpat_div(lua_State *L) { return Lpat_join(L, " / "); }

static const luaL_Reg match_metamethods[] = {{"__tostring", Lmatch_tostring}, {NULL, NULL}};

static const luaL_Reg pat_methods[] = {
    {"match", Lmatch}, {"replace", Lreplace}, {"matches", Lmatches}, {"getsource", Lpat_source}, {NULL, NULL}};

static const luaL_Reg pat_metamethods[] = {{"__gc", Lpat_gc},   {"__concat", Lpat_concat},
                                           {"__div", Lpat_div}, {"__tostring", Lpat_tostring},
                                           {"__index", NULL}, // placeholder for pat_methods
                                           {NULL, NULL}};

static const luaL_Reg bp_methods[] = {
    {"match", Lmatch}, {"replace", Lreplace}, {"compile", Lcompile}, {"matches", Lmatches}, {NULL, NULL}};

public
LUALIB_API int luaopen_bp(lua_State *L) {
    maybe_pat_t maybe_pat = bp_pattern(builtins_source, builtins_source + strlen(builtins_source));
    if (!maybe_pat.success) {
        raise_parse_error(L, maybe_pat);
        return 0;
    }
    builtins = maybe_pat.value.pat;

    lua_pushlightuserdata(L, (void *)&PAT_METATABLE);
    luaL_newlib(L, pat_metamethods);
    luaL_newlib(L, pat_methods);
    lua_setfield(L, -2, "__index");
    lua_settable(L, LUA_REGISTRYINDEX);

    lua_pushlightuserdata(L, (void *)&MATCH_METATABLE);
    luaL_newlib(L, match_metamethods);
    lua_settable(L, LUA_REGISTRYINDEX);

    luaL_newlib(L, bp_methods);
    return 1;
}