erlquery

-module(erlquery).

-export([parse/1, codegen/1]).

-type erlq_module() :: binary().
-type erlq_query() :: {binary(), binary(), binary()}.
-type erlq_method() :: {binary(), binary()}.
-type erlq_methods() :: [erlq_method()].
-type erlq_value() :: erlq_module() | erlq_query() | erlq_methods().

-record(config,
        {module :: erlq_module(), query :: erlq_query(), methods :: erlq_methods()}).

-type config() :: #config{}.

-spec parse(iodata()) -> {ok, config()} | {error, any()}.
parse(S) ->
    Clauses = split_clauses(S),
    reduce_clauses(Clauses, init_config()).

-spec init_config() -> config().
init_config() ->
    #config{module = <<>>,
            query = {<<>>, <<>>, <<>>},
            methods = []}.

-spec split_clauses(iodata()) -> [{integer(), iodata()}].
split_clauses(S) ->
    Clauses0 = string:split(S, ".\n", all),
    CLs = add_line_numbers(Clauses0),
    Clauses1 = lists:filter(fun not_whitespace/1, CLs),
    Clauses2 = lists:map(fun({L, X}) -> {L, string:trim(X, both)} end, Clauses1),
    lists:reverse(Clauses2).

-spec add_line_numbers([iodata()]) -> [{integer(), iodata()}].
add_line_numbers(Lines) ->
    add_line_numbers(Lines, 1, []).

-spec add_line_numbers([iodata()], integer(), [{integer(), iodata()}]) ->
                          [{integer(), iodata()}].
add_line_numbers([], _, Acc) ->
    Acc;
add_line_numbers([H | T], Counter, Acc) ->
    N = length(string:split(H, <<"\n">>, all)),
    add_line_numbers(T, Counter + 1 + N, [{Counter, H} | Acc]).

-spec not_whitespace({integer(), iodata()}) -> boolean().
not_whitespace({_, S}) ->
    Res = re:replace(S, "[\s\n]", "", [global]),
    B = iolist_to_binary(Res),
    B =/= <<"">>.

-spec reduce_clauses([{integer(), iodata()}], config()) ->
                        {ok, config()} | {error, any()}.
reduce_clauses([], Config) ->
    {ok, Config};
reduce_clauses([H | T], Config) ->
    case reduce_clause(H, Config) of
        {error, _Err} = Err ->
            Err;
        NewConfig ->
            reduce_clauses(T, NewConfig)
    end.

-spec reduce_clause({integer(), iodata()}, config()) -> config() | {error, any()}.
reduce_clause({Line, S}, Config) ->
    case clause_type(S) of
        module ->
            M = match_module(S),
            case validate_module(M) of
                true ->
                    Config#config{module = M};
                false ->
                    {error, {invalid_module_clause, {Line, S}}}
            end;
        query ->
            Q = match_query(S),
            case validate_query(Q) of
                true ->
                    case Q of
                        {_, _, <<"2">>} ->
                            Config#config{query = Q};
                        _ ->
                            {error, {invalid_query_arity, {Line, S}}}
                    end;
                false ->
                    {error, {invalid_query_clause, {Line, S}}}
            end;
        method ->
            Method = match_method(S),
            case validate_method(Method) of
                true ->
                    NewMethods = [Method | Config#config.methods],
                    Config#config{methods = NewMethods};
                false ->
                    {error, {invalid_method_clause, {Line, S}}}
            end;
        _ ->
            {error, {invalid_clause, {Line, S}}}
    end.

-spec clause_type(iodata()) -> module | query | method | nomatch.
clause_type(S) ->
    case is_module_clause(S) of
        true ->
            module;
        false ->
            case is_query_clause(S) of
                true ->
                    query;
                false ->
                    case is_method_clause(S) of
                        true ->
                            method;
                        false ->
                            nomatch
                    end
            end
    end.

-spec is_module_clause(iodata()) -> boolean().
is_module_clause(S) ->
    case string:prefix(S, "-module(") of
        nomatch ->
            false;
        _ ->
            true
    end.

-spec is_query_clause(iodata()) -> boolean().
is_query_clause(S) ->
    case string:prefix(S, "-query(") of
        nomatch ->
            false;
        _ ->
            true
    end.

-spec is_method_clause(iodata()) -> boolean().
is_method_clause(S) ->
    case re:run(S, "([a-zA-Z0-9]+)\s?+->") of
        {match, _} ->
            true;
        _ ->
            false
    end.

%% Expecting the .\n suffix to be stripped by split_clauses
-spec match_method(iodata()) -> erlq_method() | nomatch.
match_method(S) ->
    case re:run(S, "([a-zA-Z0-9_]+)[\s\n]?+(->)", [{capture, [1, 2]}]) of
        {match, [{Start0, End0}, {Start1, End1}]} ->
            Name = string:slice(S, Start0, End0),
            Query =
                string:trim(
                    string:slice(S, Start1 + End1), both),
            case {Name, Query} of
                {[], _} ->
                    nomatch;
                {_, []} ->
                    nomatch;
                _ ->
                    {bcast(Name), bcast(Query)}
            end;
        _ ->
            nomatch
    end.

-spec match_module(iodata()) -> erlq_module() | nomatch.
match_module(S) ->
    case re:run(S, "-module\\(([a-zA-Z0-9_]+)\\)", [{capture, [1]}]) of
        {match, [{Start, End}]} ->
            case string:slice(S, Start, End) of
                [] ->
                    nomatch;
                Slice ->
                    bcast(Slice)
            end;
        _ ->
            nomatch
    end.

%% Expecting query to look like -query(pgo:query/2).
-spec match_query(iodata()) -> erlq_query() | nomatch.
match_query(S) ->
    case re:run(S,
                "-query\\(([a-zA-Z0-9_]+):([a-zA-Z0-9_]+)/([0-9]+)\\)",
                [{capture, [1, 2, 3]}])
    of
        {match, [{Start0, End0}, {Start1, End1}, {Start2, End2}]} ->
            Mod = string:slice(S, Start0, End0),
            Fun = string:slice(S, Start1, End1),
            Arity = string:slice(S, Start2, End2),
            case {Mod, Fun, Arity} of
                {[], _, _} ->
                    nomatch;
                {_, [], _} ->
                    nomatch;
                {_, _, []} ->
                    nomatch;
                _ ->
                    {bcast(Mod), bcast(Fun), bcast(Arity)}
            end;
        _ ->
            nomatch
    end.

-spec bcast(iolist() | binary()) -> binary().
bcast(<<S/binary>>) ->
    S;
bcast(S = [_ | _]) ->
    list_to_binary(S).

-spec codegen(config()) -> {ok, binary()} | {error, {atom(), erlq_value()}}.
codegen(Config) ->
    case validate_config(Config) of
        ok ->
            M = codegen_module(Config#config.module),
            Q = codegen_query(Config#config.query),
            E = codegen_export(Config#config.methods),
            Ms = codegen_methods(Q, Config#config.methods),
            {ok, list_to_binary([M, E, Ms])};
        E ->
            E
    end.

-spec codegen_module(binary()) -> binary().
codegen_module(Module) ->
    list_to_binary([<<"-module(">>, Module, <<").\n\n">>]).

-spec codegen_query({binary(), binary(), binary()}) -> binary().
codegen_query({Module, Method, _Arity}) ->
    list_to_binary([Module, <<":">>, Method]).

-spec codegen_methods(binary(), [{binary(), binary()}]) -> binary().
codegen_methods(Query, Methods) ->
    list_to_binary(lists:map(fun(M) -> codegen_method(Query, M) end, Methods)).

-spec codegen_method(binary(), {binary(), binary()}) -> binary().
codegen_method(Query, {Method, QueryText}) ->
    list_to_binary([Method,
                    <<"(Args) ->\n">>,
                    <<"    ">>,
                    Query,
                    <<"(<<\"">>,
                    QueryText,
                    <<"\">>, Args).\n\n">>]).

-spec codegen_export([{binary(), binary()}]) -> binary().
codegen_export(Methods) ->
    list_to_binary([<<"-export([">>, reduce_exports(Methods, []), <<"\n]).\n\n">>]).

-spec reduce_exports([{binary(), binary()}], [binary()]) -> [binary()].
reduce_exports([H], Acc) ->
    lists:reverse([codegen_export_method(H) | Acc]);
reduce_exports([H | T], Acc) ->
    Res = codegen_export_method(H),
    Res2 = list_to_binary([Res, <<",\n">>]),
    reduce_exports(T, [Res2 | Acc]).

-spec codegen_export_method({binary(), binary()}) -> binary().
codegen_export_method({Name, _QueryText}) ->
    list_to_binary([Name, <<"/1">>]).

-spec validate_config(config()) -> ok | {error, {atom(), any()}}.
validate_config(Config) ->
    case validate_module(Config#config.module) of
        true ->
            case validate_query(Config#config.query) of
                true ->
                    case Config#config.query of
                        {_, _, <<"2">>} ->
                            case validate_methods(Config#config.methods) of
                                true ->
                                    ok;
                                false ->
                                    {error, {invalid_methods, Config#config.methods}}
                            end;
                        _ ->
                            {error, {invalid_query_arity, Config#config.query}}
                    end;
                false ->
                    {error, {invalid_query, Config#config.query}}
            end;
        false ->
            {error, {invalid_module, Config#config.module}}
    end.

-spec validate_module(erlq_module() | any()) -> boolean().
validate_module(<<_/binary>>) ->
    true;
validate_module(_) ->
    false.

-spec validate_query(erlq_query() | any()) -> boolean().
validate_query({<<_/binary>>, <<_/binary>>, <<_/binary>>}) ->
    true;
validate_query(_) ->
    false.

-spec validate_methods(erlq_methods() | any()) -> boolean().
validate_methods(_ = []) ->
    true;
validate_methods(_ = [H | T]) ->
    case validate_method(H) of
        true ->
            validate_methods(T);
        false ->
            false
    end.

-spec validate_method(erlq_method() | any()) -> boolean().
validate_method(_ = {<<_/binary>>, <<_/binary>>}) ->
    true;
validate_method(_) ->
    false.

-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

match_module_test() ->
    ?assertEqual(<<"foobar">>, match_module("-module(foobar).")),
    ?assertEqual(<<"foo_bar">>, match_module("-module(foo_bar).")).

match_query_test() ->
    ?assertEqual({<<"pgo">>, <<"query">>, <<"2">>}, match_query("-query(pgo:query/2)")),
    ?assertEqual({<<"pgo">>, <<"super_query">>, <<"2">>},
                 match_query("-query(pgo:super_query/2)")),
    ?assertEqual({<<"p_go">>, <<"super_query">>, <<"2">>},
                 match_query("-query(p_go:super_query/2)")).

match_method_test() ->
    ?assertEqual({<<"foobar">>, <<"SELECT * FROM foobar">>},
                 match_method("foobar ->\n    SELECT * FROM foobar")),

    ?assertEqual({<<"foobar">>, <<"SELECT * FROM foobar">>},
                 match_method("foobar->\n    SELECT * FROM foobar")),

    ?assertEqual({<<"foobar">>,
                  <<"SELECT * FROM foobar\nINNER JOIN accounts ON foobar.id = accounts.id">>},
                 match_method("foobar ->\n    SELECT * FROM foobar\nINNER JOIN accounts ON foobar.id = accounts.id")),

    ?assertEqual({<<"foo_bar">>, <<"SELECT * FROM foobar">>},
                 match_method("foo_bar ->\n    SELECT * FROM foobar")).

is_module_clause_test() ->
    ?assert(is_module_clause("-module(foobar).")),
    ?assertNot(is_module_clause("foobar")).

is_query_clause_test() ->
    ?assert(is_query_clause("-query(foobar:goo/3).")),
    ?assertNot(is_query_clause("foobar")).

is_method_clause_test() ->
    ?assert(is_method_clause("foobar -> jhgjgj")),
    ?assert(is_method_clause("foobar-> jhgjgj")),
    ?assertNot(is_method_clause("foobar")),
    ?assert(is_method_clause("foo->bar")).

clause_type_test() ->
    ?assertEqual(module, clause_type("-module(foobar)")),
    ?assertEqual(query, clause_type("-query(foobar:goo/2)")),
    ?assertEqual(method, clause_type("foo -> SELECT * FROM foobar")),
    ?assertEqual(nomatch, clause_type("foobar")).

reduce_clause_test() ->
    ?assertEqual(#config{module = <<"foo">>}, reduce_clause({0, "-module(foo)"}, #config{})),
    ?assertEqual(#config{query = {<<"foo">>, <<"bar">>, <<"2">>}},
                 reduce_clause({0, "-query(foo:bar/2)"}, #config{})),
    ?assertEqual(#config{methods = [{<<"foobar">>, <<"SELECT * FROM foobars">>}]},
                 reduce_clause({0, "foobar -> SELECT * FROM foobars"}, #config{methods = []})),

    ?assertMatch({error, {invalid_clause, {_, _}}}, reduce_clause({0, "stuff"}, #config{})),

    ?assertMatch({error, {invalid_module_clause, {_, _}}},
                 reduce_clause({0, "-module()"}, #config{})),

    ?assertMatch({error, {invalid_module_clause, {_, _}}},
                 reduce_clause({0, "-module(   )"}, #config{})),

    ?assertMatch({error, {invalid_query_arity, {_, _}}},
                 reduce_clause({0, "-query(foo:bar/3)"}, #config{})),

    ?assertMatch({error, {invalid_query_clause, {_, _}}},
                 reduce_clause({0, "-query(foo: /3)"}, #config{})),

    ?assertMatch({error, {invalid_query_clause, {_, _}}},
                 reduce_clause({0, "-query()"}, #config{})),

    ?assertMatch({error, {invalid_method_clause, {_, _}}},
                 reduce_clause({0, "foobar ->"}, #config{})),

    ?assertMatch({error, {invalid_clause, {_, _}}},
                 reduce_clause({0, " -> foobar"}, #config{})).

add_line_numbers_test() ->
    ?assertEqual([], add_line_numbers([])),
    ?assertEqual([{1, "foobar"}], add_line_numbers(["foobar"])),
    ?assertEqual([{5, "foobar"}, {1, "foobar\n\nhello"}],
                 add_line_numbers(["foobar\n\nhello", "foobar"])).

not_whitespace_test() ->
    ?assert(not_whitespace({0, "foobar"})),
    ?assertNot(not_whitespace({0, <<"">>})),
    ?assert(not_whitespace({0, <<"foobar sdfsdfsdf">>})).

split_clauses_test() ->
    S = "-module(foobar).\n\n-query(pgo:query/2).\n\nbarfoo->\n    SELECT * FROM foobars.\n",
    ?assertMatch([{_, "-module(foobar)"},
                  {_, "-query(pgo:query/2)"},
                  {_, "barfoo->\n    SELECT * FROM foobars"}],
                 split_clauses(S)).

parse_test() ->
    ?assertEqual({ok,
                  #config{module = <<"fooq">>,
                          query = {<<"foo">>, <<"bar">>, <<"2">>},
                          methods = [{<<"foobar">>, <<"SELECT * FROM foobars">>}]}},
                 parse(list_to_binary(["-module(fooq).\n\n",
                                       "-query(foo:bar/2).\n\n",
                                       "foobar ->\n",
                                       "    SELECT * FROM foobars"]))).

codegen_export_method_test() ->
    ?assertEqual(<<"foobar/1">>, codegen_export_method({"foobar", "barfoo"})).

codegen_export_test() ->
    ?assertEqual(<<"-export([foobar/1,\nbarfoo/1\n]).\n\n">>,
                 codegen_export([{"foobar", "blah"}, {"barfoo", "blahblah"}])).

codegen_query_test() ->
    ?assertEqual(<<"barfoo:query">>, codegen_query({<<"barfoo">>, <<"query">>, <<"2">>})).

codegen_module_test() ->
    ?assertEqual(<<"-module(foobar).\n\n">>, codegen_module(<<"foobar">>)).

codegen_method_test() ->
    ?assertEqual(<<"foobar(Args) ->\n    barfoo:query(<<\"SELECT * FROM foobars\">>, Args).\n\n">>,
                 codegen_method(<<"barfoo:query">>, {<<"foobar">>, <<"SELECT * FROM foobars">>})).

codegen_methods_test() ->
    ?assertEqual(<<"foobar(Args) ->\n    barfoo:query(<<\"SELECT * FROM foobars\">>, Args).\n\n">>,
                 codegen_methods(<<"barfoo:query">>,
                                 [{<<"foobar">>, <<"SELECT * FROM foobars">>}])).

codegen_test() ->
    ?assertEqual({ok,
                  <<"-module(things).\n\n-export([foobar/1\n]).\n\nfoobar(Args) ->\n    bar:foo(<<\"SELECT * FROM foobars\">>, Args).\n\n">>},
                 codegen(#config{module = <<"things">>,
                                 query = {<<"bar">>, <<"foo">>, <<"2">>},
                                 methods = [{<<"foobar">>, <<"SELECT * FROM foobars">>}]})).

compile_test() ->
    {ok, Bin} =
        codegen(#config{module = <<"testme">>,
                        query = {<<"erlquery_mock">>, <<"query">>, <<"2">>},
                        methods = [{<<"foobar">>, <<"SELECT * FROM foobars">>}]}),
    file:write_file("testme.erl", Bin),
    {ok, testme} = compile:file("testme.erl"),
    ?assertEqual(ok, testme:foobar([])),
    file:delete("testme.erl").

validate_module_test() ->
    ?assert(validate_module(<<"foo">>)),
    ?assertNot(validate_module("foo")),
    ?assertNot(validate_module(1)),
    ?assertNot(validate_module(1.0)),
    ?assertNot(validate_module(true)),
    ?assertNot(validate_module([])),
    ?assertNot(validate_module(#{})).

validate_query_test() ->
    ?assert(validate_query({<<"foo">>, <<"bar">>, <<"boo">>})),
    ?assertNot(validate_query({"foo", "bar", "boo"})),
    ?assertNot(validate_query({1, 2, 3})),
    ?assertNot(validate_query({1.0, 2.3, 4.5})).

validate_method_test() ->
    ?assert(validate_method({<<"foo">>, <<"bar">>})),
    ?assertNot(validate_method({"foo", "bar"})),
    ?assertNot(validate_method({1, 2})),
    ?assertNot(validate_method(nomatch)).

validate_methods_test() ->
    ?assert(validate_methods([])),
    ?assert(validate_methods([{<<"foo">>, <<"bar">>}])),
    ?assert(validate_methods([{<<"foo">>, <<"bar">>},
                              {<<"foo">>, <<"bar">>},
                              {<<"foo">>, <<"bar">>}])),
    ?assertNot(validate_methods([{<<"foo">>, <<"bar">>}, {"foo", "bar"}, {1, 2}])).

validate_config_test() ->
    ?assertEqual(ok,
                 validate_config(#config{module = <<"foobar">>,
                                         query = {<<"foob">>, <<"bar">>, <<"2">>},
                                         methods = [{<<"hello">>, <<"yep">>}]})),

    ?assertMatch({error, {invalid_module, _}},
                 validate_config(#config{module = 42,
                                         query = {<<"foo">>, <<"bar">>, <<"foobar">>},
                                         methods = [{<<"stuff">>, <<"hello world">>}]})),

    ?assertMatch({error, {invalid_query, _}},
                 validate_config(#config{module = <<"foobar">>,
                                         query = {1, 2, 3},
                                         methods = [{<<"stuff">>, <<"hello">>}]})),

    ?assertMatch({error, {invalid_query_arity, _}},
                 validate_config(#config{module = <<"foobar">>,
                                         query = {<<"foo">>, <<"bar">>, <<"3">>},
                                         methods = [{<<"stuff">>, <<"hello">>}]})),

    ?assertMatch({error, {invalid_methods, _}},
                 validate_config(#config{module = <<"foobar">>,
                                         query = {<<"foo">>, <<"bar">>, <<"2">>},
                                         methods = [{1, 2}]})).

-endif.