%%%-------------------------------------------------------------------
%%% Licensed to the Apache Software Foundation (ASF) under one
%%% or more contributor license agreements.  See the NOTICE file
%%% distributed with this work for additional information
%%% regarding copyright ownership.  The ASF licenses this file
%%% to you under the Apache License, Version 2.0 (the
%%% "License"); you may not use this file except in compliance
%%% with the License.  You may obtain a copy of the License at
%%%
%%%   http://www.apache.org/licenses/LICENSE-2.0
%%%
%%% Unless required by applicable law or agreed to in writing,
%%% software distributed under the License is distributed on an
%%% "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
%%% KIND, either express or implied.  See the License for the
%%% specific language governing permissions and limitations
%%% under the License.
%%%

-module(grpc_lib_compile).

-export([file/2]).

%%% ============================================================================
%%% Exported functions
%%% ============================================================================

file(Filename, Options) ->
    {Generate, OtherOptions} = grpc_lib:keytake(generate, Options, server),
    GPBModule = compile_pb(Filename, OtherOptions),
    case Generate of 
        client ->
            write_client(GPBModule);
        server ->
            write_server(GPBModule)
    end.

%%% ============================================================================
%%% Internal functions
%%% ============================================================================

compile_pb(Filename, Options) ->
    ok = gpb_compile:file(Filename, [maps | Options ++ [{i, "."}]]),
    CompiledPB =  filename:rootname(Filename) ++ ".erl",
    GpbInludeDir = filename:join(code:lib_dir(gpb), "include"),
    {ok, Module, Compiled} = compile:file(CompiledPB, 
                                          [binary, {i, GpbInludeDir}]),
    {module, _} = code:load_binary(Module, CompiledPB, Compiled),
    Module.

write_server(Module) ->
    ModuleName = atom_to_list(Module) ++ "_server",
    Header = print_header(ModuleName),
    Exports = print_server_exports(Module),
    Messages = print_messages(Module),
    Decoder = print_decoder(Module),
    Services = print_services(Module, server),
    ok = file:write_file(ModuleName ++ ".erl", [Header, Exports, Messages,
                                                Decoder, Services]).

write_client(Module) ->
    ModuleName = atom_to_list(Module) ++ "_client",
    Header = print_header(ModuleName),
    Exports = print_client_exports(Module),
    TypeExports = print_type_exports(Module),
    Messages = print_messages(Module),
    Decoder = print_decoder(Module),
    Services = print_services(Module, client),
    ok = file:write_file(ModuleName ++ ".erl", [Header, Exports, Messages,
                                                TypeExports, Decoder,
                                                Services]).

print_header(ModuleName) ->
    ["-module(", string_to_atom(ModuleName), ").\n\n"
     "%% this file was generated by grpc\n\n"].

print_server_exports(Module) ->
    RPCs = rpcs(Module),
    to_io_list(["-export([",
                separate(["decoder/0"] ++
                          [[Name, "/3"]|| #{name := Name} <- RPCs], ",\n         "),
                "]).\n\n"]).

%% Only "simple rpcs" (no streaming) are considered.
print_client_exports(Module) ->
    RPCs = rpcs(Module),
    to_io_list(["-export([",
                separate([[Name, "/3"]|| #{name := Name,
                                           output_stream := false,
                                           input_stream := false} <- RPCs], ",\n         "),
                "]).\n\n"]).

%% Export all types, to avoid warnings.
print_type_exports(Module) ->
    to_io_list(["-export_type([",
                separate([print_type_export(M) || M <- Module:get_msg_defs()],
                         ",\n              "),
                "]).\n\n"]).

rpcs(Module) ->
    lists:flatten([begin
                       {_, Rs} = Module:get_service_def(S),
                       Rs
                   end || S <- Module:get_service_names()]).
 
print_messages(Module) -> 
    to_io_list([print_message(M) || M <- Module:get_msg_defs()]).

print_decoder(Module) ->
    to_io_list(["-spec decoder() -> module().\n",
                "%% The module (generated by gpb) used to encode and decode protobuf\n"
                "%% messages.\n",
                "decoder() -> ", Module, ".\n\n"]).

print_services(Module, Role) ->
    to_io_list([print_service(Module:get_service_def(S), Role) 
                || S <- Module:get_service_names()]).

print_service({{service, Name}, RPCs}, Role) ->
    ["%% RPCs for service ", Name, "\n\n",
     print_rpcs(RPCs, Name, Role)].

print_rpcs(RPCs, _Servcie, server) ->
    [print_rpc(R) || R <- RPCs];
print_rpcs(RPCs, Service, client) ->
    [print_rpc_call(R, Service) || #{input_stream := false,
                                     output_stream := false} = R <- RPCs].

print_rpc(#{input_stream := InStream, output_stream := OutStream} = Rpc) ->
    case {InStream, OutStream} of 
        {false, false} ->
            print_unary_rpc(Rpc);
        {false, true} -> 
            print_server_to_client_streaming_rpc(Rpc);
        {true, false} ->
            print_client_to_server_streaming_rpc(Rpc);
        {true, true} ->
            print_bidirectional_streaming_rpc(Rpc)
    end.
 
print_unary_rpc(#{name := Name,  
                  output := Output} = Rpc) ->
    [print_spec_call(Rpc),
     "    {", Output, "(), grpc:stream()} | grpc:error_response().\n",
     "%% This is a unary RPC\n",
     Name, "(_Message, Stream, _State) ->\n",
     "    {#{}, Stream}.\n\n"].

print_server_to_client_streaming_rpc(#{name := Name, 
                                       output := Output} = Rpc) ->
    [print_spec_call(Rpc),
     "    {[", Output, "()], grpc:stream()} | grpc:error_response().\n",
     "%% This is a server-to-client streaming RPC\n",
     Name, "(_Message, Stream, _State) ->\n",
     "    {[#{}], Stream}.\n\n"].

print_client_to_server_streaming_rpc(#{name := Name, 
                                       input := Input, 
                                       output := Output}) ->
    ["-spec ", Name, 
     "(Message::", Input, "() | eof, Stream::grpc:stream(), State::any()) ->\n",
     "    {continue, grpc:stream(), any()} |\n",
     "    {", Output, "(), grpc:stream()} |\n",
     "    grpc:error_response().\n",
     "%% This is a client-to-server streaming RPC. After the client has sent the last message\n",
     "%% this function will be called a final time with 'eof' as the first argument. This last\n",
     "%% invocation must return the response message.\n",
     Name, "(_Message, Stream, State) ->\n",
     "    {continue, Stream, State}.\n\n"].

print_bidirectional_streaming_rpc(#{name := Name, 
                                    input := Input, 
                                    output := Output}) ->
    ["-spec ", Name, 
     "(Message::", Input, "() | eof, Stream::grpc:stream(), State::any()) ->\n",
     "    {[", Output, "()], grpc:stream(), any()} |\n",
     "    {[", Output, "()], grpc:stream()} |\n",
     "    grpc:error_response().\n",
     "%% This is a bidirectional streaming RPC. If the client terminates the stream\n",
     "%% this function will be called a final time with 'eof' as the first argument.\n", 
     Name, "(_Message, Stream, State) ->\n",
     "    {[], Stream, State}.\n\n"].
 
%% the client side, only for non-streaming rpcs.
print_rpc_call(#{name := Name,  
                 input := Input,
                 output := Output} = _Rpc, Service) ->
    ["-spec ", Name, "(\n",
     "        Connection::grpc_client:connection(),\n", 
     "        Message::", Input, "(),\n",
     "        Options::[grpc_client:stream_option() |\n",
     "                  {timeout, timeout()}]) ->\n",
     "        grpc_client:unary_response(", Output, "()).\n",
     "%% This is a unary RPC\n",
     Name, "(Connection, Message, Options) ->\n",
     "    grpc_client:unary(Connection, Message,\n"
     "                      ", Service, ", ", Name, ",\n",
     "                       decoder(), Options).\n\n"].

print_spec_call(#{name := Name, input := Input}) ->
    ["-spec ", Name, 
     "(Message::", Input, "(), Stream::grpc:stream(), State::any()) ->\n"].

print_message({{msg, Name}, Fields}) -> 
    ["-type ", Name, "() ::\n    #{", print_fields(Fields), "}.\n\n"];
print_message({{enum, Name}, Options}) -> 
    ["-type ", Name, "() ::\n    ", print_options(Options), ".\n\n"].

print_type_export({{msg, Name}, _Fields}) -> 
    [Name, "/0"];
print_type_export({{enum, Name}, _Values}) -> 
    [Name, "/0"].

print_fields(Fields) -> 
    separate([print_field(F) || F <- Fields], ",\n      ").

print_options(Options) -> 
    separate([V || {V, _} <- Options], " |\n    ").

print_field(#{name := Name, type := Type, occurrence := Occurrence}) ->
    [Name, 
     case {Occurrence, Type} of 
         {optional, _} -> " => ";
         {required, _} -> " := ";
         %% Maps are marked as 'repeated' by gpb, but in fact they 
         %% have 1 value (a single map, not a list).
         {repeated, {map, _, _}}  -> " => ";
         {repeated, _}  -> " => ["
     end, 
     print_type(Type),
     case {Occurrence, Type} of 
         {repeated, {map, _, _}}  -> "";
         {repeated, _}  -> "]";
         _ -> ""
     end];
print_field(#{name := Name, fields := Alternatives}) ->
    [Name, " =>\n          ",
     separate([print_alternative(F) || F <- Alternatives], " |\n          ")].

print_alternative(#{name := Name, type := Type}) ->
    ["{", Name, ", ", print_type(Type), "}"].

print_type(int32) -> "integer()";
print_type(int64) -> "integer()";
print_type(uint32) -> "integer()";
print_type(uint64) -> "integer()";
print_type(sint32) -> "integer()";
print_type(sint64) -> "integer()";
print_type(fixed32) -> "integer()";
print_type(fixed64) -> "integer()";
print_type(sfixed32) -> "integer()";
print_type(sfixed64) -> "integer()";
print_type(string) -> "string()";
print_type(bool) -> "boolean()";
print_type(bytes) -> "binary()";
print_type(double) -> "float() | infinity | '-infinity' | nan";
print_type(float) -> "float() | infinity | '-infinity' | nan";
print_type({enum, Enum}) -> [Enum, "() | integer()"];
print_type({map, KeyType, ValueType}) -> 
    ["#{", print_type(KeyType), " => ", print_type(ValueType), "}"];
print_type({msg, MsgName}) -> [MsgName, "()"].

separate([], _) ->
    [];
separate(L, With) ->
    [Last | Rest] = lists:reverse(L),
    separate(Rest, With, [Last]).

separate([], _, Acc) ->
    Acc;
separate([I | T], With, Acc) ->
    separate(T, With, [I, With | Acc]).


to_io_list(L) ->
    print2(lists:flatten(L)).

print2(L) ->
    [lists:flatten([print3(C) || C <- L])].

print3(C) when is_atom(C) ->
    io_lib:format("~p", [C]);
print3(C) ->
    C.

string_to_atom(S) when is_list(S) ->
    io_lib:format("~p", [list_to_atom(S)]).