-module(grpcbox_stream).

-include_lib("chatterbox/include/http2.hrl").
-include_lib("kernel/include/logger.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,
         ctx/1,
         ctx/2,
         handle_streams/2,
         handle_call/2,
         handle_info/2]).

-export([init/3,
         on_receive_headers/2,
         on_send_push_promise/2,
         on_receive_data/2,
         on_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(),
                services_table      :: ets:tid(),
                req_headers=[]      :: list(),
                full_method         :: binary() | undefined,
                input_ref           :: reference() | undefined,
                callback_pid        :: pid() | undefined,
                connection_pid      :: pid(),
                request_encoding    :: gzip | identity | undefined,
                response_encoding   :: gzip | identity | 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,
                stats_handler       :: module() | undefined,
                stats               :: term() | undefined}).

-type 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()}.

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

on_receive_headers(Headers, State=#state{ctx=_Ctx}) ->
    %% proplists:get_value(<<":method">>, Headers) =:= <<"POST">>,
    Metadata = grpcbox_utils:headers_to_metadata(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,

    FullPath = proplists:get_value(<<":path">>, Headers),
    %% wait to rpc_begin here since we need to know the method
    Ctx1 = ctx:with_value(Ctx, grpc_server_method, FullPath),
    State1=#state{ctx=Ctx2} = stats_handler(Ctx1, rpc_begin, {}, State#state{full_method=FullPath}),

    RequestEncoding = parse_options(<<"grpc-encoding">>, Headers),
    ResponseEncoding = parse_options(<<"grpc-accept-encoding">>, Headers),
    ContentType = parse_options(<<"content-type">>, Headers),

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

    handle_service_lookup(Ctx2, string:lexemes(FullPath, "/"),
                          State1#state{resp_headers=RespHeaders,
                                       req_headers=Headers,
                                       request_encoding=RequestEncoding,
                                       response_encoding=ResponseEncoding,
                                       content_type=ContentType}).

handle_service_lookup(Ctx, [Service, Method], State=#state{services_table=ServicesTable}) ->
    case ets:lookup(ServicesTable, {Service, Method}) of
        [M=#method{}] ->
            State1 = State#state{ctx=Ctx,
                                 method=M},
            handle_auth(Ctx, State1);
        _ ->
            end_stream(?GRPC_STATUS_UNIMPLEMENTED, <<"Method not found on server">>, State)
    end;
handle_service_lookup(_, _, State) ->
    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}.

handle_auth(_Ctx, State=#state{auth_fun=AuthFun,
                               socket=Socket,
                               method=#method{input={_, InputStreaming}}}) ->
    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, State#state{handler=self()}]),
                    {ok, State#state{input_ref=Ref,
                                     callback_pid=Pid}};
                _ ->
                    {ok, State}
            end;
        _ ->
            end_stream(?GRPC_STATUS_UNAUTHENTICATED, <<"">>, State)
    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_stream => false},
                  StreamInterceptor(Ref, State, ServerInfo, fun Module:Function/2)
          end) of
        {ok, Response, State2} ->
            send(Response, State2);
        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_stream => 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_data(_, State=#state{method=undefined}) ->
    {ok, State};
on_receive_data(Bin, State=#state{request_encoding=Encoding,
                                  buffer=Buffer}) ->
    try
        {NewBuffer, Messages} = grpcbox_frame:split(<<Buffer/binary, Bin/binary>>, Encoding),
        State1 = lists:foldl(fun(EncodedMessage, StateAcc=#state{}) ->
                                     StateAcc1 = handle_message(EncodedMessage, StateAcc),
                                     StateAcc1
                             end, State, Messages),
        {ok, State1#state{buffer=NewBuffer}}
    catch
        throw:{grpc_error, {Status, Message}} ->
            end_stream(Status, Message, State);
        C:E:S ->
            ?LOG_INFO("crash: class=~p exception=~p stacktrace=~p", [C, E, S]),
            end_stream(?GRPC_STATUS_UNKNOWN, <<>>, State)
    end.

handle_message(EncodedMessage, State=#state{input_ref=Ref,
                                            ctx=Ctx,
                                            callback_pid=Pid,
                                            method=#method{proto=Proto,
                                                           input={Input, InputStream},
                                                           output={_Output, OutputStream}}}) ->
    try Proto:decode_msg(EncodedMessage, Input) of
        Message ->
            State1=#state{ctx=Ctx1} =
                stats_handler(Ctx, in_payload, #{uncompressed_size => erlang:external_size(Message),
                                                 compressed_size => size(EncodedMessage)}, State),
            case {InputStream, OutputStream} of
                {true, _} ->
                    Pid ! {Ref, Message},
                    State1;
                {false, true} ->
                    _ = proc_lib:spawn_link(?MODULE, handle_streams,
                                            [Message, State1#state{handler=self()}]),
                    State1;
                {false, false} ->
                    handle_unary(Ctx1, Message, State1)
            end
    catch
        error:{gpb_error, _} ->
            ?THROW(?GRPC_STATUS_INTERNAL, <<"Error parsing request protocol buffer">>)
    end.

handle_unary(Ctx, Message, State=#state{unary_interceptor=UnaryInterceptor,
                                        full_method=FullMethod,
                                        method=#method{module=Module,
                                                       function=Function,
                                                       proto=_Proto,
                                                       input={_Input, _InputStream},
                                                       output={_Output, _OutputStream}}}) ->
    Ctx1 = ctx_with_stream(Ctx, State),
    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} ->
            State1 = from_ctx(Ctx2),
            send(false, Response, State1);
        E={grpc_error, _} ->
            throw(E)
    end.

on_end_stream(State) ->
    on_end_stream_(State),
    {ok, State}.

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

%% Internal

stats_handler(Ctx, _, _, State=#state{stats_handler=undefined}) ->
    State#state{ctx=Ctx};
stats_handler(Ctx, Event, Stats, State=#state{stats_handler=StatsHandler,
                                              stats=StatsState}) ->
    {Ctx1, StatsState1} = StatsHandler:handle(Ctx, server, Event, Stats, StatsState),
    State#state{ctx=Ctx1,
                stats=StatsState1}.

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,
                                         ctx=Ctx,
                                         resp_trailers=Trailers}) ->
    EncodedTrailers = grpcbox_utils:encode_headers(Trailers),
    h2_connection:send_trailers(ConnPid, StreamId, [{<<"grpc-status">>, Status},
                                                    {<<"grpc-message">>, Message} | EncodedTrailers],
                                [{send_end_stream, true}]),
    Ctx1 = ctx:with_value(Ctx, grpc_server_status, grpcbox_utils:status_to_string(Status)),
    State1 = stats_handler(Ctx1, rpc_end, {}, State),
    {ok, State1#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(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 = grpcbox_utils: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)).

ctx(#state{handler=Pid}) ->
    h2_stream:call(Pid, ctx).

ctx(#state{handler=Pid}, Ctx) ->
    h2_stream:call(Pid, {ctx, Ctx}).

handle_call(ctx, State=#state{ctx=Ctx}) ->
    {ok, Ctx, State};
handle_call({ctx, Ctx}, State) ->
    {ok, ok, State#state{ctx=Ctx}}.

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;
handle_info(_, 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{ctx=Ctx,
                                connection_pid=ConnPid,
                                stream_id=StreamId,
                                response_encoding=Encoding,
                                method=#method{proto=Proto,
                                               input={_Input, _},
                                               output={Output, _}}}) ->
    BodyToSend = Proto:encode_msg(Message, Output),
    OutFrame = grpcbox_frame:encode(Encoding, BodyToSend),
    ok = h2_connection:send_body(ConnPid, StreamId, OutFrame, [{send_end_stream, End}]),
    stats_handler(Ctx, out_payload, #{uncompressed_size => erlang:external_size(Message),
                                      compressed_size => size(BodyToSend)}, State).

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

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 ->
            identity;
        <<"gzip", _/binary>> ->
            gzip;
        <<"snappy", _/binary>> ->
            snappy;
        <<"deflate", _/binary>> ->
            deflate;
        <<"identity", _/binary>> ->
            identity;
        _ ->
            ?THROW(?GRPC_STATUS_UNIMPLEMENTED, <<"unknown encoding">>)
    end.

maybe_encode_headers(Headers) ->
    maps:map(fun(K, V) ->
                     maybe_encode_header_value(K, V)
             end, Headers).

maybe_encode_header_value(K, V) ->
    case binary:longest_common_suffix([K, <<"-bin">>]) == 4 of
        true ->
            base64:encode(V);
        false ->
            V
    end.