-module(grpcbox_stream).

-include_lib("chatterbox/include/http2.hrl").
-include("grpcbox.hrl").

-behaviour(h2_stream).

-export([send/2,
         send/3,
         send_headers/2,
         add_headers/2,
         add_trailers/2,
         set_trailers/2,
         code_to_status/1,
         error/2,
         handle_streams/2,
         handle_info/2]).

-export([init/3,
         on_receive_request_headers/2,
         on_send_push_promise/2,
         on_receive_request_data/2,
         on_request_end_stream/1]).

-export_type([t/0,
              grpc_status/0,
              grpc_status_message/0,
              grpc_error/0,
              grpc_error_response/0]).

-record(state, {handler            :: pid(),
                socket,
                auth_fun,
                buffer             :: binary(),
                ctx                :: ctx:ctx(),
                req_headers=[]     :: list(),
                full_method        :: binary() | undefined,
                input_ref          :: reference() | undefined,
                callback_pid       :: pid() | undefined,
                connection_pid     :: pid(),
                request_encoding   :: gzip | undefined,
                response_encoding  :: gzip | undefined,
                content_type       :: proto | json | undefined,
                resp_headers=[]    :: list(),
                resp_trailers=[]   :: list(),
                headers_sent=false :: boolean(),
                trailers_sent=false :: boolean(),
                unary_interceptor  :: fun() | undefined,
                stream_interceptor :: fun() | undefined,
                stream_id          :: stream_id(),
                method             :: #method{} | undefined}).

-opaque t() :: #state{}.

-type grpc_status_message() :: unicode:chardata().
-type grpc_status() :: 0..16.
-type grpc_error() :: {grpc_status(), grpc_status_message()}.
-type grpc_error_response() :: {grpc_error, grpc_error()}.

-define(GRPC_ERROR(Status, Message), {grpc_error, {Status, Message}}).
-define(THROW(Status, Message), throw(?GRPC_ERROR(Status, Message))).

init(ConnPid, StreamId, [Socket, AuthFun, UnaryInterceptor, StreamInterceptor]) ->
    process_flag(trap_exit, true),
    {ok, #state{connection_pid=ConnPid,
                stream_id=StreamId,
                buffer = <<>>,
                auth_fun=AuthFun,
                unary_interceptor=UnaryInterceptor,
                stream_interceptor=StreamInterceptor,
                socket=Socket,
                handler=self()}}.

on_receive_request_headers(Headers, State=#state{auth_fun=AuthFun,
                                                 socket=Socket}) ->
    %% proplists:get_value(<<":method">>, Headers) =:= <<"POST">>,
    RequestEncoding = parse_options(<<"grpc-encoding">>, Headers),
    ResponseEncoding = parse_options(<<"grpc-accept-encoding">>, Headers),
    ContentType = parse_options(<<"content-type">>, Headers),

    Metadata =  lists:foldl(fun({K, V}, Acc) ->
                                    case is_reserved_header(K) of
                                        true ->
                                            Acc;
                                        false ->
                                            maps:put(K, maybe_decode_header(K, V), Acc)
                                    end
                            end, #{}, Headers),

    Ctx = case parse_options(<<"grpc-timeout">>, Headers) of
              infinity ->
                  grpcbox_metadata:new_incoming_ctx(Metadata);
              D ->
                  ctx:with_deadline_after(grpcbox_metadata:new_incoming_ctx(Metadata), D, nanosecond)
          end,

    RespHeaders = [{<<":status">>, <<"200">>},
                   {<<"user-agent">>, <<"grpc-erlang/0.1.0">>},
                   {<<"content-type">>, content_type(ContentType)}
                   | response_encoding(ResponseEncoding)],

    FullPath = proplists:get_value(<<":path">>, Headers),
    case string:lexemes(FullPath, "/") of
        [Service, Method] ->
            case ets:lookup(?SERVICES_TAB, {Service, Method}) of
                [M=#method{input={_, InputStreaming}}] ->
                    State1 = State#state{resp_headers=RespHeaders,
                                         req_headers=Headers,
                                         full_method=FullPath,
                                         request_encoding=RequestEncoding,
                                         response_encoding=ResponseEncoding,
                                         content_type=ContentType,
                                         ctx=Ctx,
                                         method=M},
                    case authenticate(sock:peercert(Socket), AuthFun) of
                        {true, _Identity} ->
                            case InputStreaming of
                                true ->
                                    Ref = make_ref(),
                                    Pid = proc_lib:spawn_link(?MODULE, handle_streams,
                                                              [Ref, State1#state{handler=self()}]),
                                    {ok, State1#state{input_ref=Ref,
                                                      callback_pid=Pid}};
                                _ ->
                                    {ok, State1}
                            end;
                        _ ->
                            end_stream(?GRPC_STATUS_UNAUTHENTICATED, <<"">>, State1)
                    end;
                _ ->
                    end_stream(?GRPC_STATUS_UNIMPLEMENTED, <<"Method not found on server">>,
                               State#state{resp_headers=RespHeaders})
            end;
        _ ->
            State1 = State#state{resp_headers=[{<<":status">>, <<"200">>},
                                               {<<"user-agent">>, <<"grpc-erlang/0.1.0">>}]},
            end_stream(?GRPC_STATUS_UNIMPLEMENTED, <<"failed parsing path">>, State1),
            {ok, State1}
    end.

authenticate(_, undefined) ->
    {true, undefined};
authenticate({ok, Cert}, Fun) ->
    Fun(Cert);
authenticate(_, _) ->
    false.

handle_streams(Ref, State=#state{full_method=FullMethod,
                                 stream_interceptor=StreamInterceptor,
                                 method=#method{module=Module,
                                                function=Function,
                                                output={_, false}}}) ->
    case (case StreamInterceptor of
              undefined -> Module:Function(Ref, State);
              _ ->
                  ServerInfo = #{full_method => FullMethod,
                                 service => Module,
                                 input_stream => true,
                                 output_strema => false},
                  StreamInterceptor(Ref, State, ServerInfo, fun Module:Function/2)
          end) of
        {ok, Response, State1} ->
            send(false, Response, State1);
        E={grpc_error, _} ->
            throw(E)
    end;
handle_streams(Ref, State=#state{full_method=FullMethod,
                                 stream_interceptor=StreamInterceptor,
                                 method=#method{module=Module,
                                                function=Function,
                                                output={_, true}}}) ->
    case StreamInterceptor of
        undefined ->
            Module:Function(Ref, State);
        _ ->
            ServerInfo = #{full_method => FullMethod,
                           service => Module,
                           input_stream => true,
                           output_strema => true},
            StreamInterceptor(Ref, State, ServerInfo, fun Module:Function/2)
    end.

on_send_push_promise(_, State) ->
    {ok, State}.

ctx_with_stream(Ctx, Stream) ->
    ctx:set(Ctx, ctx_stream_key, Stream).

from_ctx(Ctx) ->
    ctx:get(Ctx, ctx_stream_key).

on_receive_request_data(Bin, State=#state{request_encoding=Encoding,
                                          input_ref=Ref,
                                          callback_pid=Pid,
                                          unary_interceptor=UnaryInterceptor,
                                          ctx=Ctx,
                                          buffer=Buffer,
                                          full_method=FullMethod,
                                          method=#method{module=Module,
                                                         function=Function,
                                                         proto=Proto,
                                                         input={Input, InputStream},
                                                         output={_Output, OutputStream}}})->
    try
        {NewBuffer, Messages} = split_frame(<<Buffer/binary, Bin/binary>>, Encoding),
        State1 = lists:foldl(fun(EncodedMessage, StateAcc) ->
                                     try Proto:decode_msg(EncodedMessage, Input) of
                                         Message ->
                                             case {InputStream, OutputStream} of
                                                 {true, _} ->
                                                     Pid ! {Ref, Message},
                                                     StateAcc;
                                                 {false, true} ->
                                                     _ = proc_lib:spawn_link(?MODULE, handle_streams,
                                                                             [Message, StateAcc#state{handler=self()}]),
                                                     StateAcc;
                                                 {false, false} ->
                                                     Ctx1 = ctx_with_stream(Ctx, StateAcc),
                                                     case (case UnaryInterceptor of
                                                               undefined -> Module:Function(Ctx1, Message);
                                                               _ ->
                                                                   ServerInfo = #{full_method => FullMethod,
                                                                                  service => Module},
                                                                   UnaryInterceptor(Ctx1, Message, ServerInfo, fun Module:Function/2)
                                                           end) of
                                                         {ok, Response, Ctx2} ->
                                                             StateAcc1 = from_ctx(Ctx2),
                                                             send(false, Response, StateAcc1);
                                                         E={grpc_error, _} ->
                                                             throw(E)
                                                     end
                                             end
                                     catch
                                         error:{gpb_error, _} ->
                                             ?THROW(?GRPC_STATUS_INTERNAL, <<"Error parsing request protocol buffer">>)
                                     end
                             end, State, Messages),
        {ok, State1#state{buffer=NewBuffer}}
    catch
        throw:{grpc_error, {Status, Message}} ->
            end_stream(Status, Message, State);
        _C:_T ->
            end_stream(?GRPC_STATUS_UNKNOWN, <<>>, State)
    end.

on_request_end_stream(State=#state{input_ref=Ref,
                                   callback_pid=Pid,
                                   method=#method{input={_Input, true},
                                                  output={_Output, false}}}) ->
    Pid ! {Ref, eos},
    {ok, State};
on_request_end_stream(State=#state{input_ref=Ref,
                                   callback_pid=Pid,
                                   method=#method{input={_Input, true},
                                                  output={_Output, true}}}) ->
    Pid ! {Ref, eos},
    {ok, State};
on_request_end_stream(State=#state{input_ref=_Ref,
                                   callback_pid=_Pid,
                                   method=#method{input={_Input, false},
                                                  output={_Output, true}}}) ->
    {ok, State};
on_request_end_stream(State=#state{method=#method{output={_Output, false}}}) ->
    end_stream(State),
    {ok, State};
on_request_end_stream(State) ->
    end_stream(State),
    {ok, State}.

%% Internal

end_stream(State) ->
    end_stream(?GRPC_STATUS_OK, <<>>, State).

end_stream(Status, Message, State=#state{headers_sent=false}) ->
    end_stream(Status, Message, send_headers(State));
end_stream(_Status, _Message, State=#state{trailers_sent=true}) ->
    {ok, State};
end_stream(Status, Message, State=#state{connection_pid=ConnPid,
                                         stream_id=StreamId,
                                         resp_trailers=Trailers}) ->
    EncodedTrailers = encode_headers(Trailers),
    h2_connection:send_trailers(ConnPid, StreamId, [{<<"grpc-status">>, Status},
                                                    {<<"grpc-message">>, Message} | EncodedTrailers],
                                [{send_end_stream, true}]),
    {ok, State#state{trailers_sent=true}}.

set_trailers(Ctx=#{}, Trailers) ->
    State = from_ctx(Ctx),
    ctx_with_stream(Ctx, State#state{resp_trailers=maps:to_list(Trailers)}).

send_headers(State) ->
    send_headers([], State).

send_headers(Ctx=#{}, Headers) when is_map(Headers) ->
    State = from_ctx(Ctx),
    send_headers(maps:to_list(grpc_lib:maybe_encode_headers(Headers)), State);

send_headers(_Metadata, State=#state{headers_sent=true}) ->
    State;
send_headers(Metadata, State=#state{connection_pid=ConnPid,
                                    stream_id=StreamId,
                                    resp_headers=Headers,
                                    headers_sent=false}) ->
    MdHeaders = encode_headers(Metadata),
    h2_connection:send_headers(ConnPid, StreamId, Headers++MdHeaders, [{send_end_stream, false}]),
    State#state{headers_sent=true}.

code_to_status(0) -> ?GRPC_STATUS_OK;
code_to_status(1) -> ?GRPC_STATUS_CANCELLED;
code_to_status(2) -> ?GRPC_STATUS_UNKNOWN;
code_to_status(3) -> ?GRPC_STATUS_INVALID_ARGUMENT;
code_to_status(4) -> ?GRPC_STATUS_DEADLINE_EXCEEDED;
code_to_status(5) -> ?GRPC_STATUS_NOT_FOUND;
code_to_status(6) -> ?GRPC_STATUS_ALREADY_EXISTS;
code_to_status(7) -> ?GRPC_STATUS_PERMISSION_DENIED;
code_to_status(8) -> ?GRPC_STATUS_RESOURCE_EXHAUSTED;
code_to_status(9) -> ?GRPC_STATUS_FAILED_PRECONDITION;
code_to_status(10) -> ?GRPC_STATUS_ABORTED;
code_to_status(11) -> ?GRPC_STATUS_OUT_OF_RANGE;
code_to_status(12) -> ?GRPC_STATUS_UNIMPLEMENTED;
code_to_status(13) -> ?GRPC_STATUS_INTERNAL;
code_to_status(14) -> ?GRPC_STATUS_UNAVAILABLE;
code_to_status(15) -> ?GRPC_STATUS_DATA_LOSS;
code_to_status(16) -> ?GRPC_STATUS_UNAUTHENTICATED.

error(Status, Message) ->
    exit(?GRPC_ERROR(Status, Message)).

handle_info({add_headers, Headers}, State) ->
    update_headers(Headers, State);
handle_info({add_trailers, Trailers}, State) ->
    update_trailers(Trailers, State);
handle_info({send_proto, Message}, State) ->
    send(false, Message, State);
handle_info({'EXIT', _, normal}, State) ->
    end_stream(State),
    State;
handle_info({'EXIT', _, {grpc_error, {Status, Message}}}, State) ->
    end_stream(Status, Message, State),
    State;
handle_info({'EXIT', _, _Other}, State) ->
    end_stream(?GRPC_STATUS_UNKNOWN, <<"process exited without reason">>, State),
    State.

add_headers(Headers, #state{handler=Pid}) ->
    Pid ! {add_headers, Headers}.

add_trailers(Ctx=#{}, Trailers=#{}) ->
    State=#state{resp_trailers=RespTrailers} = from_ctx(Ctx),
    ctx_with_stream(Ctx, State#state{resp_trailers=maps:to_list(Trailers)++RespTrailers});
add_trailers(Headers, #state{handler=Pid}) ->
    Pid ! {add_trailers, Headers}.

update_headers(Headers, State=#state{resp_headers=RespHeaders}) ->
    State#state{resp_headers=RespHeaders++Headers}.

update_trailers(Trailers, State=#state{resp_trailers=RespTrailers}) ->
    State#state{resp_trailers=RespTrailers++Trailers}.

send(Message, #state{handler=Pid}) ->
    Pid ! {send_proto, Message}.

send(End, Message, State=#state{headers_sent=false}) ->
    State1 = send_headers(State),
    send(End, Message, State1);
send(End, Message, State=#state{connection_pid=ConnPid,
                                stream_id=StreamId,
                                response_encoding=Encoding,
                                method=#method{proto=Proto,
                                               input={_Input, _},
                                               output={Output, _}}}) ->
    BodyToSend = Proto:encode_msg(Message, Output),
    OutFrame = encode_frame(Encoding, BodyToSend),
    ok = h2_connection:send_body(ConnPid, StreamId, OutFrame, [{send_end_stream, End}]),
    State.

encode_frame(gzip, Bin) ->
    CompressedBin = zlib:gzip(Bin),
    Length = byte_size(CompressedBin),
    <<1, Length:32, CompressedBin/binary>>;
encode_frame(_, Bin) ->
    Length = byte_size(Bin),
    <<0, Length:32, Bin/binary>>.

split_frame(Frame, Encoding) ->
    split_frame(Frame, Encoding, []).

split_frame(<<>>, _Encoding, Acc) ->
    {<<>>, lists:reverse(Acc)};
split_frame(<<0, Length:32, Encoded:Length/binary, Rest/binary>>, Encoding, Acc) ->
    split_frame(Rest, Encoding, [Encoded | Acc]);
split_frame(<<1, Length:32, Compressed:Length/binary, Rest/binary>>, Encoding, Acc) ->
    Encoded = case Encoding of
                  gzip ->
                      try zlib:gunzip(Compressed)
                      catch
                          error:data_error ->
                              ?THROW(?GRPC_STATUS_INTERNAL,
                                     <<"Could not decompress but compression algorithm supported">>)
                      end;
                  _ ->
                      ?THROW(?GRPC_STATUS_UNIMPLEMENTED,
                             <<"Compression mechanism used by client not supported by the server">>)
              end,
    split_frame(Rest, Encoding, [Encoded | Acc]);
split_frame(Bin, _Encoding, Acc) ->
    {Bin, lists:reverse(Acc)}.

response_encoding(gzip) ->
    [{<<"grpc-encoding">>, <<"gzip">>}];
response_encoding(snappy) ->
    [{<<"grpc-encoding">>, <<"snappy">>}];
response_encoding(deflate) ->
    [{<<"grpc-encoding">>, <<"deflate">>}];
response_encoding(undefined) ->
    [].

content_type(json) ->
    <<"application/grpc+json">>;
content_type(_) ->
    <<"application/grpc+proto">>.

timeout_to_duration(T, <<"H">>) ->
    erlang:convert_time_unit(timer:hours(T), millisecond, nanosecond);
timeout_to_duration(T, <<"M">>) ->
    erlang:convert_time_unit(timer:minutes(T), millisecond, nanosecond);
timeout_to_duration(T, <<"S">>) ->
    erlang:convert_time_unit(T, second, nanosecond);
timeout_to_duration(T, <<"m">>) ->
    erlang:convert_time_unit(T, millisecond, nanosecond);
timeout_to_duration(T, <<"u">>) ->
    erlang:convert_time_unit(T, microsecond, nanosecond);
timeout_to_duration(T, <<"n">>) ->
    timer:seconds(T).

parse_options(<<"grpc-timeout">>, Headers) ->
    case proplists:get_value(<<"grpc-timeout">>, Headers, infinity) of
        infinity ->
            infinity;
        T ->
            {I, U} = string:to_integer(T),
            timeout_to_duration(I, U)
    end;
parse_options(<<"content-type">>, Headers) ->
    case proplists:get_value(<<"content-type">>, Headers, undefined) of
        undefined ->
            proto;
        <<"application/grpc">> ->
            proto;
        <<"application/grpc+proto">> ->
            proto;
        <<"application/grpc+json">> ->
            json;
        <<"application/grpc+", _>> ->
            ?THROW(?GRPC_STATUS_UNIMPLEMENTED, <<"unknown content-type">>)
    end;
parse_options(<<"grpc-encoding">>, Headers) ->
    parse_encoding(<<"grpc-encoding">>, Headers);
parse_options(<<"grpc-accept-encoding">>, Headers) ->
    parse_encoding(<<"grpc-accept-encoding">>, Headers).

parse_encoding(EncodingType, Headers) ->
    case proplists:get_value(EncodingType, Headers, undefined) of
        undefined ->
            undefined;
        <<"gzip">> ->
            gzip;
        <<"snappy">> ->
            snappy;
        <<"deflate">> ->
            deflate;
        _ ->
            ?THROW(?GRPC_STATUS_UNIMPLEMENTED, <<"unknown encoding">>)
    end.


%% TODO: consolidate with grpc_lib. But have to update their header map to support
%% a list of values for a key.

maybe_decode_header(Key, Value) ->
    case binary:longest_common_suffix([Key, <<"-bin">>]) == 4 of
        true ->
            decode_header(Value);
        false ->
            Value
    end.

%% golang gRPC implementation does not add the padding that the Erlang
%% decoder needs...
decode_header(Base64) when byte_size(Base64) rem 4 == 3 ->
    base64:decode(<<Base64/bytes, "=">>);
decode_header(Base64) when byte_size(Base64) rem 4 == 2 ->
    base64:decode(<<Base64/bytes, "==">>);
decode_header(Base64) ->
    base64:decode(Base64).

encode_headers([]) ->
    [];
encode_headers([{Key, Value} | Rest]) ->
     case binary:longest_common_suffix([Key, <<"-bin">>]) == 4 of
         true ->
             [{Key, base64:encode(Value)} | encode_headers(Rest)];
         false ->
             [{Key, Value} | encode_headers(Rest)]
     end.

is_reserved_header(<<"content-type">>) -> true;
is_reserved_header(<<"grpc-message-type">>) -> true;
is_reserved_header(<<"grpc-encoding">>) -> true;
is_reserved_header(<<"grpc-message">>) -> true;
is_reserved_header(<<"grpc-status">>) -> true;
is_reserved_header(<<"grpc-timeout">>) -> true;
is_reserved_header(<<"grpc-status-details-bin">>) -> true;
is_reserved_header(<<"te">>) -> true;
is_reserved_header(_) -> false.