-module(tpFlame).

-export([
   pfs/2             %% 分析单个文件
   , pfm/2           %% 分析多个文件
]).

-record(state, {
   outputPath = "",
   pid,
   lastTs,
   count = 0,
   acc = []
}).

-spec pfs(file:filename_all(), file:filename_all()) -> ok.
pfs(InputFile, OutputPath) ->
   {ok, FinalState} = tpFReader:fold(fun handleEvent/2, #state{outputPath = OutputPath}, InputFile),
   flush(FinalState).

-spec pfm(file:filename(), file:filename()) -> ok.
pfm(InputFiles, OutputPath) ->
   PfFiles = filelib:wildcard(InputFiles),
   doPfm(PfFiles, #state{outputPath = OutputPath}).

doPfm([], State) ->
   flush(State);
doPfm([InputFile | PfFiles], State) ->
   {ok, NewState} = tpFReader:fold(fun handleEvent/2, State, InputFile),
   doPfm(PfFiles, NewState).

handleEvent({Type, Pid, Ts, Arg}, State) ->
   doExp({trace_ts, Pid, Type, Arg, Ts}, State);
handleEvent({Type, Pid, Ts, Arg, ExtraOrMspec}, State) ->
   doExp({trace_ts, Pid, Type, Arg, ExtraOrMspec, Ts}, State);
handleEvent({Type, Pid, Ts, Arg, Extra, Mspec}, State) ->
   doExp({trace_ts, Pid, Type, Arg, Extra, Mspec, Ts}, State).

doExp(T, #state{outputPath = OutputPath} = State) ->
   trace_ts = element(1, T),
   Pid = element(2, T),
   PidState =
      case erlang:get(Pid) of
         undefined ->
            io:format("~p ", [Pid]),
            #state{outputPath = OutputPath};
         SomeState ->
            SomeState
      end,
   NewPidState = doExpInner(T, PidState),
   erlang:put(Pid, NewPidState),
   State.

%% in & out, without call context, don't help us
doExpInner({trace_ts, _Pid, InOut, _MFA, _TS}, #state{lastTs = undefined} = PS) when InOut == in; InOut == out ->
   PS;
%% return_from and return_to, without call context, don't help us
doExpInner({trace_ts, _Pid, Return, _MFA, _TS}, #state{lastTs = undefined} = PS) when Return == return_from; Return == return_to ->
   PS;
doExpInner({trace_ts, Pid, call, MFA, BIN, TS}, #state{lastTs = LastTS, acc = Acc, count = Count} = PS) ->
   try
      %% Calculate time elapsed, TS-LastTs.
      %% 0. If Acc is empty, then skip step #1.
      %% 1. Credit elapsed time to the stack on the top of Acc.
      %% 2. Push a 0 usec item with this stack onto Acc.
      Stak =
         lists:filter(
            fun(<<"unknown function">>) -> false;
               (_) -> true
            end, stak_binify(BIN)),
      Stack0 = stak_trim(Stak),
      MFA_bin = mfa_binify(MFA),
      Stack1 = [MFA_bin | lists:reverse(Stack0)],
      Acc2 =
         case Acc of
            [] ->
               [{Stack1, 0}];
            [{LastStack, LastTime} | Tail] ->
               USec = TS - LastTS,
%                   io:format("Stack1: ~p ~p\n", [Stack1, USec]),
               [{Stack1, 0},
                  {LastStack, LastTime + USec} | Tail]
         end,
      %% TODO: more state tracking here.
      PS#state{pid = Pid, lastTs = TS, count = Count + 1, acc = Acc2}
   catch Class:Reason:StackTrace ->
      io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, Class, Reason, StackTrace]),
      PS
   end;
doExpInner({trace_ts, _Pid, return_to, MFA, TS}, #state{lastTs = LastTS, acc = Acc} = S) ->
   try
      %% Calculate time elapsed, TS-LastTs.
      %% 1. Credit elapsed time to the stack on the top of Acc.
      %% 2. Push a 0 usec item with the "best" stack onto Acc.
      %%    "best" = MFA exists in the middle of the stack onto Acc,
      %%    or else MFA exists at the top of a stack elsewhere in Acc.
      [{LastStack, LastTime} | Tail] = Acc,
      MFA_bin = mfa_binify(MFA),
      BestStack = lists:dropwhile(fun(SomeMFA) when SomeMFA /= MFA_bin -> true;
         (_) -> false
                                  end, find_matching_stack(MFA_bin, Acc)),
      USec = TS - LastTS,
      Acc2 = [{BestStack, 0},
         {LastStack, LastTime + USec} | Tail],
%    io:format(user, "return-to: ~p\n", [lists:sublist(Acc2, 4)]),
      S#state{lastTs = TS, acc = Acc2}
   catch XX:YY:ZZ ->
      io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, XX, YY, ZZ]),
      S
   end;

doExpInner({trace_ts, _Pid, gc_start, _Info, TS}, #state{lastTs = LastTS, acc = Acc} = S) ->
   try
      %% Push a 0 usec item onto Acc.
      [{LastStack, LastTime} | Tail] = Acc,
      NewStack = [<<"GARBAGE-COLLECTION">> | LastStack],
      USec = TS - LastTS,
      Acc2 = [{NewStack, 0},
         {LastStack, LastTime + USec} | Tail],
%    io:format(user, "GC 1: ~p\n", [lists:sublist(Acc2, 4)]),
      S#state{lastTs = TS, acc = Acc2}
   catch _XX:_YY:_ZZ ->
      %% io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, _XX, _YY, _ZZ]),
      S
   end;
doExpInner({trace_ts, _Pid, gc_end, _Info, TS}, #state{lastTs = LastTS, acc = Acc} = S) ->
   try
      %% Push the GC time onto Acc, then push 0 usec item from last exec
      %% stack onto Acc.
      [{GCStack, GCTime}, {LastExecStack, _} | Tail] = Acc,
      USec = TS - LastTS,
      Acc2 = [{LastExecStack, 0}, {GCStack, GCTime + USec} | Tail],
%    io:format(user, "GC 2: ~p\n", [lists:sublist(Acc2, 4)]),
      S#state{lastTs = TS, acc = Acc2}
   catch _XX:_YY:_ZZ ->
      %% io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, _XX, _YY, _ZZ]),
      S
   end;

doExpInner({trace_ts, _Pid, out, MFA, TS}, #state{lastTs = LastTS, acc = Acc} = S) ->
   try
      %% Push a 0 usec item onto Acc.
      %% The MFA reported here probably doesn't appear in the stacktrace
      %% given to us by the last 'call', so add it here.
      [{LastStack, LastTime} | Tail] = Acc,
      MFA_bin = mfa_binify(MFA),
      NewStack = [<<"SLEEP">>, MFA_bin | LastStack],
      USec = TS - LastTS,
      Acc2 = [{NewStack, 0},
         {LastStack, LastTime + USec} | Tail],
      S#state{lastTs = TS, acc = Acc2}
   catch XX:YY:ZZ ->
      io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, XX, YY, ZZ]),
      S
   end;
doExpInner({trace_ts, _Pid, in, MFA, TS}, #state{lastTs = LastTS, acc = Acc} = S) ->
   try
      %% Push the Sleep time onto Acc, then push 0 usec item from last
      %% exec stack onto Acc.
      %% The MFA reported here probably doesn't appear in the stacktrace
      %% given to us by the last 'call', so add it here.
      MFA_bin = mfa_binify(MFA),
      [{SleepStack, SleepTime}, {LastExecStack, _} | Tail] = Acc,
      USec = TS - LastTS,
      Acc2 = [{[MFA_bin | LastExecStack], 0}, {SleepStack, SleepTime + USec} | Tail],
      S#state{lastTs = TS, acc = Acc2}
   catch XX:YY:ZZ ->
      io:format(user, "~p: ~p:~p @ ~p\n", [?LINE, XX, YY, ZZ]),
      S
   end;

%exp1_inner(end_of_trace = _Else, #state{pid=Pid, output_path=OutputPath, acc=Acc} = S) ->
%    {ok, FH} = file:open(OutputPath, [write, raw, binary, delayed_write]),
%    io:format("Writing to ~s ... ", [OutputPath]),
%    [begin
%         Pid_str = io_lib:format("~w", [Pid]),
%         Time_str = integer_to_list(Time),
%         file:write(FH, [Pid_str, $;, intersperse($;, lists:reverse(Stack)), 32, Time_str, 10])
%     end || {Stack, Time} <- Acc],
%    file:close(FH),
%    io:format("finished\n"),
%    S;
doExpInner(_Else, S) ->
%    io:format("?? ~P\n", [_Else, 10]),
   S.

find_matching_stack(MFA_bin, [{H, _Time} | _] = Acc) ->
   case lists:member(MFA_bin, H) of
      true ->
         H;
      false ->
         find_matching_stack2(MFA_bin, Acc)
   end.

find_matching_stack2(MFA_bin, [{[MFA_bin | _StackTail] = Stack, _Time} | _]) ->
   Stack;
find_matching_stack2(MFA_bin, [_H | T]) ->
   find_matching_stack2(MFA_bin, T);
find_matching_stack2(_MFA_bin, []) ->
   [<<"FIND-MATCHING-FAILED">>].

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

intersperse(_, []) -> [];
intersperse(_, [X]) -> [X];
intersperse(Sep, [X | Xs]) -> [X, Sep | intersperse(Sep, Xs)].

stak_trim([<<"proc_lib:init_p_do_apply/3">>, <<"gen_fsm:decode_msg/9">>, <<"gen_fsm:handle_msg/7">>, <<"gen_fsm:loop/7">> | T]) ->
   stak_trim([<<"GEN-FSM">> | T]);
stak_trim([<<"GEN-FSM">>, <<"gen_fsm:decode_msg/9">>, <<"gen_fsm:handle_msg/7">>, <<"gen_fsm:loop/7">> | T]) ->
   stak_trim([<<"GEN-FSM">> | T]);
stak_trim(Else) ->
   Else.

stak_binify(Bin) when is_binary(Bin) ->
   [list_to_binary(X) || X <- stak(Bin)];
stak_binify(X) ->
   list_to_binary(io_lib:format("~w", [X])).

mfa_binify({M, F, A}) ->
   list_to_binary(io_lib:format("~w:~w/~w", [M, F, A]));
mfa_binify(X) ->
   list_to_binary(io_lib:format("~w", [X])).

%% Borrowed from redbug.erl

stak(Bin) ->
   lists:foldl(fun munge/2, [], string:tokens(binary_to_list(Bin), "\n")).

munge(I, Out) ->
   case I of %% lists:reverse(I) of
      "..." ++ _ -> ["truncated!!!" | Out];
      _ ->
         case string:str(I, "Return addr") of
            0 ->
               case string:str(I, "cp = ") of
                  0 -> Out;
                  _ -> [mfaf(I) | Out]
               end;
            _ ->
               case string:str(I, "erminate process normal") of
                  0 -> [mfaf(I) | Out];
                  _ -> Out
               end
         end
   end.

mfaf(I) ->
   [_, C | _] = string:tokens(I, "()+"),
   string:strip(C).

flush(#state{outputPath = OutputPath}) ->
   PidStates = get(),
   {ok, FH} = file:open(OutputPath, [write, raw, binary, delayed_write]),
   io:format("\n\nWriting to ~s for ~w processes... ", [OutputPath, length(PidStates)]),
   _ = [
      [begin
          Pid_str0 = lists:flatten(io_lib:format("~w", [Pid])),
          Size = length(Pid_str0),
          Pid_str = [$(, lists:sublist(Pid_str0, 2, Size - 2), $)],
          Time_str = integer_to_list(Time),
          file:write(FH, [Pid_str, $;, intersperse($;, lists:reverse(Stack)), 32, Time_str, 10])
       end || {Stack, Time} <- Acc]
      || {Pid, #state{acc = Acc} = _S} <- PidStates],
   _ = file:close(FH),
   io:format("finished!\n"),
   ok.