%%% -*- erlang -*-
%%%
%%% This file is part of hackney released under the Apache 2 license.
%%% See the NOTICE for more information.


-module(hackney_manager).
-behaviour(gen_server).

-export([new_request/1,
         start_async_response/1,
         stop_async_response/1,
         cancel_request/1,
         close_request/1,
         controlling_process/2]).

-export([get_state/1, get_state/2,
         update_state/1, update_state/2,
         store_state/1, store_state/2,
         take_control/2,
         handle_error/1]).

-export([async_response_pid/1,
         with_async_response_pid/2]).

-export([start_link/0]).


%% private gen_server api
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
         terminate/2, code_change/3]).

-include("hackney.hrl").

-record(request, {ref,
                  pid,
                  async_pid=nil,
                  state}).

-record(request_info, {pool,
                       start_time,
                       host}).


-define(REFS, hackney_manager_refs).

-record(mstate, {pids=dict:new(),
                 mod_metrics}).


new_request(#client{request_ref=Ref}=Client) when is_reference(Ref) ->
    {ok, StartTime} = take_control(Ref, Client),
     {Ref, Client#client{start_time=StartTime}};
new_request(Client) ->
    {Ref, StartTime} = init_request(Client),
    {Ref, Client#client{start_time=StartTime, request_ref=Ref}}.


init_request(InitialState) ->
    %% initialize the request
    Ref = make_ref(),
    %% store the current state in the process dictionnary
    put(Ref, InitialState#client{request_ref=Ref}),
    %% supervise the process owner
    {ok, StartTime} = gen_server:call(?MODULE, {new_request, self(), Ref,
                                                InitialState}, infinity),
    {Ref, StartTime}.


cancel_request(#client{request_ref=Ref}) ->
    cancel_request(Ref);
cancel_request(Ref) when is_reference(Ref) ->
    case get_state(Ref) of
        req_not_found ->
            req_not_found;
        #client{socket=Skt}=Client when Skt /= nil ->
            #client{transport=Transport, socket=Socket, buffer=Buffer,
                    response_state=RespState} = Client,

            %% only the owner can cancel the request
            case Transport:controlling_process(Socket, self()) of
                ok ->
                    %% remove the request
                    erase(Ref),

                    %% stop to monitor the request
                    case gen_server:call(?MODULE, {cancel_request, Ref}) of
                        ok ->
                            %% return the latest state
                            {ok, {Transport, Socket, Buffer, RespState}};
                        Error ->
                            Error
                    end;
                Error ->
                    Error
            end;
        Client ->
            #client{transport=Transport, socket=Socket,
                    buffer=Buffer, response_state=RespState} = Client,

            %% remove the request
            erase(Ref),


            %% stop to monitor the request
            case gen_server:call(?MODULE, {cancel_request, Ref}) of
                ok ->
                    %% return the latest state
                    {ok, {Transport, Socket, Buffer, RespState}};
                Error ->
                    Error
            end
    end.

close_request(#client{}=Client) ->
    #client{transport=Transport,
            socket=Socket,
            state=Status,
            request_ref=Ref} = Client,


    %% remove the request
    erase(Ref),
    ets:delete(?MODULE, Ref),

    %% stop to monitor the request
    Reply = gen_server:call(?MODULE, {cancel_request, Ref}),

    case Status of
        done -> ok;
        _ when Socket /= nil ->
            catch Transport:controlling_process(Socket, self()),
            catch Transport:close(Socket),
            ok;
        _ -> ok
    end,
    Reply;
close_request(Ref) ->
    case get_state(Ref) of
        req_not_found ->
            req_not_found;
        Client ->
            close_request(Client)
    end.


controlling_process(Ref, Pid) ->
    case get(Ref) of
        undefined ->
            {error, not_owner};
        Client ->
            Reply = gen_server:call(?MODULE, {controlling_process, Ref, Pid}),
            case Reply of
                ok ->
                    #client{transport=Transport, socket=Socket} = Client,
                    Transport:controlling_process(Socket, Pid),
                    ets:insert(?MODULE, {Ref, #request{ref=Ref, state=Client}}),
                    ok;
                Error ->
                    Error
            end
    end.

start_async_response(Ref) ->
    case get_state(Ref)of
        req_not_found ->
            req_not_found;
        Client ->
            #client{transport=Transport, socket=Socket,
                    stream_to=StreamTo} = Client,
            case gen_server:call(?MODULE, {start_async_response, Ref,
                                           StreamTo, Client}) of
                {ok, Pid} ->
                    %% store temporarely the socket in the the ets so it can
                    %% be used by the other process later
                    true = ets:insert(?MODULE, {Ref, #request{ref=Ref,
                                                        state=Client}}),
                    %% delete the current state from the process dictionnary
                    %% since it's not the owner
                    erase(Ref),

                    %% transfert the control of the socket
                    Transport:controlling_process(Socket, Pid);
                Error ->
                    Error
            end
    end.

stop_async_response(Ref) ->
    gen_server:call(?MODULE, {stop_async_response, Ref, self()}, infinity).

async_response_pid(Ref) ->
    case ets:lookup(?REFS, Ref) of
        [] ->
            {error, req_not_found};
        [{Ref, {_, nil, _}}] ->
            {error, req_not_async};
        [{Ref, {_, Pid, _}}] ->
            {ok, Pid}
    end.

with_async_response_pid(Ref, Fun) ->
    case async_response_pid(Ref) of
        {ok, Pid} ->
            Fun(Pid);
        Error ->
            Error
    end.

get_state(#client{request_ref=Ref}) ->
    get_state(Ref);

get_state(Ref) ->
    case get(Ref) of
        undefined ->
            case ets:lookup(?MODULE, Ref) of
                [] ->
                    req_not_found;
                [{Ref, #request{state=State}}] ->
                    %% store the state in the new context, only the current
                    %% owner can handle it.
                    put(Ref, State),
                    %% delete the state, from ets
                    ets:delete(?MODULE, Ref),
                    State
            end;
        State ->
            State
    end.

get_state(Ref, Fun) ->
    case get_state(Ref) of
        req_not_found -> {error, req_not_found};
        State -> Fun(State)
    end.

update_state(#client{request_ref=Ref}=NState) ->
    update_state(Ref, NState).

update_state(Ref, NState) ->
    put(Ref, NState).

store_state(#client{request_ref=Ref}=NState) ->
    store_state(Ref, NState).

store_state(Ref, NState) ->
    true = ets:insert(?MODULE, {Ref, #request{ref=Ref, state=NState}}),
    ok.

take_control(Ref, NState) ->
    %% maybe delete the state from ets
    ets:delete(?MODULE, Ref),
    %% add the state to the current context
    put(Ref, NState),
    gen_server:call(?MODULE, {take_control, Ref, NState}, infinity).

handle_error(#client{request_ref=Ref, dynamic=true}) ->
    close_request(Ref);

handle_error(#client{request_ref=Ref, transport=Transport,
                    socket=Socket}=Client) ->
    case get_state(Ref) of
        req_not_found -> ok;
        _ ->
            catch Transport:controlling_process(Socket, self()),
            catch Transport:close(Socket),
            NClient = Client#client{socket=nil, state=closed},
            put(Ref, NClient),
            ok
    end.


start_link() ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).

init(_) ->
    ets:new(hackney_pool, [
                    named_table,
                    set,
                    public
            ]),

    ets:new(?MODULE, [
                    set,
                    {keypos, 1},
                    public,
                    named_table,
                    {read_concurrency, true},
                    {write_concurrency, true}
            ]),

    ets:new(?REFS, [named_table, set, protected]),

    %% initialize metrics
    Mod = init_metrics(),

    process_flag(trap_exit, true),
    %% return {ok, {Pids, Refs}}
    %% Pids are the managed pids
    %% Refs are the managed requests
    {ok, #mstate{pids=dict:new(),
                 mod_metrics=Mod}}.


handle_call({new_request, Pid, Ref, Client}, _From, #mstate{pids=Pids}=State) ->
    %% get pool name
    Pool = proplists:get_value(pool, Client#client.options),
    %% set requInfo
    StartTime = os:timestamp(),
    ReqInfo = #request_info{pool=Pool,
                            start_time=StartTime,
                            host=Client#client.host},

    %% start the request
    start_request(ReqInfo, State),

    %% link the request owner
    link(Pid),
    %% store the pid
    Pids2 = dict:store(Pid, {Ref, owner}, Pids),
    ets:insert(?REFS, {Ref, {Pid, nil, ReqInfo}}),
    {reply, {ok, StartTime}, State#mstate{pids=Pids2}};

handle_call({take_control, Ref, Client}, _From, State) ->
    StartTime = os:timestamp(),

    case ets:lookup(?REFS, Ref) of
        [] ->
            %% not supposed to happen but ignore it.
            {reply, {ok, StartTime}, State};
        [{Ref, {Owner, Stream, Info}}] ->
            NInfo = Info#request_info{start_time=StartTime,
                                      host=Client#client.host},
            %% start the request
            start_request(NInfo, State),

            ets:insert(?REFS, {Ref, {Owner, Stream, NInfo}}),
            {reply, {ok, StartTime}, State}
    end;

handle_call({start_async_response, Ref, StreamTo, Client}, _From, State) ->
    %% start the stream and eventually update the owner of the request
    case do_start_async_response(Ref, StreamTo, Client, State) of
        {ok, Pid, NState} ->
            {reply, {ok, Pid}, NState};
        Error ->
            {reply, Error, State}
    end;

handle_call({stop_async_response, Ref, To}, _From, State) ->
    case ets:lookup(?REFS, Ref) of
        [] -> {reply, {ok, Ref}, State};
        [{Ref, {_Owner, nil, _Info}}] ->
            %% there is no async request to handle, just return
            {ok, Ref};
        [{Ref, {Owner, Stream, Info}}] ->
            %% tell to the stream to stop
            Stream ! {Ref, stop_async, self()},
            receive
                {Ref, ok} ->
                    %% if the stream return, we unlink it and update the
                    %% state. if we stop the async request and want to use it
                    %% in another process, make sure to unlink the old owner
                    %% and link the new one.
                    unlink(Stream),
                    ets:insert(?REFS, {Ref, {To, nil, Info}}),
                    Pids1 = dict:erase(Stream, State#mstate.pids),

                    Pids2 = case To of
                                Owner ->
                                    %% same owner do nothing
                                    Pids1;
                                _ ->
                                    %% new owner, link it and un link the old
                                    %% one
                                    unlink(Owner),
                                    link(To),
                                    dict:store(To, {Ref, owner},
                                               dict:erase(Owner, Pids1))
                            end,

                    {reply, {ok, Ref}, State#mstate{pids=Pids2}};
                Error ->
                    {reply, Error, State}
            after 5000 ->
                      {reply, {error, timeout}, State}
            end
    end;

handle_call({controlling_process, Ref, Pid}, _From, State) ->
    case ets:lookup(?REFS, Ref) of
        [] -> {reply, badarg, State};
        [{Ref, {Pid, _, _}}] ->
            %% the request is already controlled by this process just return
            {reply, ok, State};
        [{Ref, {Owner, Stream, Info}}] ->
            %% new owner, link it.
            unlink(Owner),
            link(Pid),
            Pids2 = dict:store(Pid, {Ref, owner},
                               dict:erase(Owner, State#mstate.pids)),
            ets:insert(?REFS, {Ref, {Pid, Stream, Info}}),
            {reply, ok, State#mstate{pids=Pids2}}
    end;

handle_call({cancel_request, Ref}, _From, State) ->
    PoolHandler = hackney_app:get_app_env(pool_handler, hackney_pool),
    case ets:lookup(?REFS, Ref) of
        [] -> {reply, badarg, State};
        [{Ref, {Owner, nil, #request_info{pool=Pool}=Info}}] ->
            %% no stream just cancel the request and unlink the owner.
            unlink(Owner),
            ets:delete(?REFS, Ref),
            Pids2 = dict:erase(Owner, State#mstate.pids),
            %% notify the pool that the request have been canceled
            PoolHandler:notify(Pool, {'DOWN', Ref, request, Owner, cancel}),

            %% update metrics
            finish_request(Info, State),

            {reply, ok, State#mstate{pids=Pids2}};
        [{Ref, {Owner, Stream, #request_info{pool=Pool}=Info}}] when is_pid(Stream) ->
            unlink(Owner),
            unlink(Stream),
            Pids2 = dict:erase(Stream, dict:erase(Owner, State#mstate.pids)),
            ets:delete(?REFS, Ref),
            %% notify the pool that the request have been canceled
            PoolHandler:notify(Pool, {'DOWN', Ref, request, Owner, cancel}),

            %% update metrics
            finish_request(Info, State),

            %% terminate the async response
            case terminate_async_response(Stream) of
                ok ->
                    {reply, ok, State#mstate{pids=Pids2}};
                Error ->
                    {reply, Error, State#mstate{pids=Pids2}}
            end
    end.

handle_cast(_Msg, Children) ->
    {noreply, Children}.

handle_info({'EXIT', Pid, Reason}, State) ->
    case dict:find(Pid, State#mstate.pids) of
        {ok, PidInfo} ->
            handle_exit(Pid, PidInfo, Reason, State);
        _ ->
            {noreply, State}
    end;

handle_info(_Info, State) ->
    {noreply, State}.


code_change(_OldVsn, Ring, _Extra) ->
   {ok, Ring}.

terminate(_Reason, _State) ->
    ok.

do_start_async_response(Ref, StreamTo, Client, State) ->
    %% get current owner
    [{Ref, {Owner, _, Info}}] = ets:lookup(?REFS, Ref),

    %% if not stream target we use the owner
    StreamTo2 = case StreamTo of
                    false -> Owner;
                    _ -> StreamTo
                end,

    %% start the stream process
    case catch hackney_stream:start_link(StreamTo2, Ref, Client) of
        {ok, Pid} when is_pid(Pid) ->
            ets:insert(?REFS, {Ref, {StreamTo2, Pid, Info}}),
            Pids2 = case StreamTo2 of
                        Owner ->
                            dict:store(Pid, {Ref, stream}, State#mstate.pids);
                        _ ->
                            %% unlink and replace the old owner by the new
                            %% target of the request
                            unlink(Owner),
                            Pids1 = dict:store(StreamTo2, {Ref, stream},
                                               dict:erase(Owner,
                                                          State#mstate.pids)),
                            %% store stthe stream
                            dict:store(Pid, {Ref, stream}, Pids1)
                    end,
            {ok, Pid, State#mstate{pids=Pids2}};
        {error, What} ->
            {error, What};
        What ->
            {error, What}
    end.


%% a stream exited
handle_exit(Pid, {Ref, stream}, Reason, State) ->
    %% delete the pid from our list
    Pids1 = dict:erase(Pid, State#mstate.pids),

    case ets:lookup(?REFS, Ref) of
        [] ->
            %% ref already removed just return
            {noreply, State#mstate{pids=Pids1}};
        [{Ref, {Owner, Pid, #request_info{pool=Pool}=Info}}] ->
            %% unlink the owner
            unlink(Owner),
            Pids2 = dict:erase(Pid, Pids1),
            %% if anormal reason let the owner knows
            case Reason of
                normal ->
                    ok;
                _ ->
                    Owner ! {'DOWN', Ref, Reason}
            end,
            %% remove the reference
            ets:delete(?REFS, Ref),
            ets:delete(?MODULE, Ref),

            %% notify the pool that the request have been canceled
            PoolHandler = hackney_app:get_app_env(pool_handler, hackney_pool),
            PoolHandler:notify(Pool, {'DOWN', Ref, request, Owner, Reason}),

            %% update metrics
            finish_request(Info, State),

            %% reply
            {noreply, State#mstate{pids=Pids2}}
    end;
%% owner exited
handle_exit(Pid, {Ref, owner}, Reason, State) ->
    PoolHandler = hackney_app:get_app_env(pool_handler, hackney_pool),
    %% delete the pid from our list
    Pids1 = dict:erase(Pid, State#mstate.pids),
    case ets:lookup(?REFS, Ref) of
        [] ->
            %% ref already removed just return
            {noreply, State#mstate{pids=Pids1}};
        [{Ref, {Pid, nil, #request_info{pool=Pool}=Info}}] ->
            %% no stream
            %% remove the reference
            ets:delete(?REFS, Ref),
            ets:delete(?MODULE, Ref),
            %% notify the pool that the request have been canceled
            PoolHandler:notify(Pool, {'DOWN', Ref, request, Pid, Reason}),
            %% update metrics
            finish_request(Info, State),
            %% reply
            {noreply, State#mstate{pids=Pids1}};
        [{Ref, {Pid, Stream, #request_info{pool=Pool}=Info}}] ->
            unlink(Stream),
            Pids2 = dict:erase(Stream, Pids1),
            %% terminate the async stream
            terminate_async_response(Stream),
            %% remove the reference
            ets:delete(?REFS, Ref),
            ets:delete(?MODULE, Ref),
            %% notify the pool that the request have been canceled
            PoolHandler:notify(Pool, {'DOWN', Ref, request, Pid, Reason}),
            %% update metrics
            finish_request(Info, State),
            {noreply, State#mstate{pids=Pids2}}
    end.

monitor_child(Pid) ->
    erlang:monitor(process, Pid),
    unlink(Pid),

    receive
        {'EXIT', Pid, Reason} ->
            receive
                {'DOWN', _, process, Pid, _} ->
                    {error, Reason}
            end
    after 0 ->
            ok
    end.

terminate_async_response(Stream) ->
    terminate_async_response(Stream, shutdown).

terminate_async_response(Stream, Reason) ->
    case monitor_child(Stream) of
        ok ->
            exit(Stream, Reason),
            wait_async_response(Stream);
        Error ->
            Error
    end.

wait_async_response(Stream) ->
    receive
        {'DOWN', _MRef, process, Stream, _Reason} ->
            ok
    end.

init_metrics() ->
    %% get metrics module
    Mod = hackney_util:mod_metrics(),

    %% initialise metrics
    Mod:new(counter, [hackney, nb_requests]),
    Mod:new(counter, [hackney, total_requests]),
    Mod:new(counter, [hackney, finished_requests]),
    Mod.

start_request(#request_info{host=Host}, #mstate{mod_metrics=Mod}) ->
    Mod:increment_counter([hackney, Host, nb_requests]),
    Mod:increment_counter([hackney, nb_requests]),
    Mod:increment_counter([hackney, total_requests]).


finish_request(#request_info{start_time=Begin, host=Host},
               #mstate{mod_metrics=Mod}) ->
    RequestTime = timer:now_diff(os:timestamp(), Begin)/1000,
    Mod:update_histogram([hackney, Host, request_time], RequestTime),
    Mod:decrement_counter([hackney, Host, nb_requests]),
    Mod:decrement_counter([hackney, nb_requests]),
    Mod:increment_counter([hackney, finished_requests]).