%% @copyright 2019-2024 Guilherme Andrade
%%
%% Permission is hereby granted, free of charge, to any person obtaining a
%% copy  of this software and associated documentation files (the "Software"),
%% to deal in the Software without restriction, including without limitation
%% the rights to use, copy, modify, merge, publish, distribute, sublicense,
%% and/or sell copies of the Software, and to permit persons to whom the
%% Software is furnished to do so, subject to the following conditions:
%%
%% The above copyright notice and this permission notice shall be included in
%% all copies or substantial portions of the Software.
%%
%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
%% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
%% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
%% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
%% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
%% FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
%% DEALINGS IN THE SOFTWARE.

-module(lager_transform).

-ifdef(E48).
-moduledoc "Parse transform".
-endif.

-include_lib("stdlib/include/assert.hrl").
-include("lager.hrl").

%%-------------------------------------------------------------------
%% API Function Exports
%%-------------------------------------------------------------------

-export([parse_transform/2]).

-ignore_xref([parse_transform/2]).

%%-------------------------------------------------------------------
%% Internal API Function Exports
%%-------------------------------------------------------------------

-export([get_pr_context/1]).

%% ------------------------------------------------------------------
%% Macro Definitions
%% ------------------------------------------------------------------

-define(pr_context_attribute_name, '___$fake_lager.pr_context').

%%-------------------------------------------------------------------
%% Record and Type Definitions
%%-------------------------------------------------------------------

-record(module_context, {
    name :: module(),
    file :: string() | undefined,
    sinks :: [atom(), ...],
    pr_context :: fake_lager_pr:context()
}).

-record(function_context, {
    module :: module(),
    file :: string() | undefined,
    name :: atom(),
    arity :: arity(),
    sinks :: [atom(), ...]
}).

%% ------------------------------------------------------------------
%% Static Check Tweaks
%% ------------------------------------------------------------------

-elvis([
    {elvis_style, macro_naming_convention, disable},
    {elvis_style, no_throw, disable},
    {elvis_text_style, max_line_length, disable}
]).

%%-------------------------------------------------------------------
%% API Function Definitions
%%-------------------------------------------------------------------

-spec parse_transform(term(), [tuple()]) -> term().
parse_transform(Ast, Options) ->
    _ = check_for_unsupported_options(Options),
    ExtraSinks = proplists:get_value(lager_extra_sinks, Options, []),
    Sinks = [lager | ExtraSinks],

    %write_terms("ast_before.txt", Ast),
    Module = get_module(Ast),
    File = get_file(Ast),

    InitialContext = #module_context{
        name = Module,
        file = File,
        sinks = Sinks,
        pr_context = fake_lager_pr:new_context()
    },
    {MappedAst, FinalContext} =
        lists:mapfoldl(
            fun mapfold_ast_statement/2,
            InitialContext,
            Ast
        ),

    AstWithRecordDefs =
        insert_pr_context_attribute(
            MappedAst,
            FinalContext#module_context.pr_context
        ),

    %write_terms("ast_after.txt", AstWithRecordDefs),
    AstWithRecordDefs.

%%-------------------------------------------------------------------
%% Internal API Function Definitions
%%-------------------------------------------------------------------

-spec get_pr_context(Module) -> Context when
    Module :: module(),
    Context :: fake_lager_pr:context().
%% @private
get_pr_context(Module) ->
    Attributes = apply(Module, module_info, [attributes]),

    case lists:keyfind(?pr_context_attribute_name, 1, Attributes) of
        {?pr_context_attribute_name, [PrContext]} ->
            PrContext;
        false ->
            throw(no_pr_context)
    end.

%%-------------------------------------------------------------------
%% Internal Function Definitions - Metadata
%%-------------------------------------------------------------------

get_module(Ast) ->
    ModuleAttribute = lists:keyfind(module, 3, Ast),
    ?assertNotEqual(false, ModuleAttribute),

    case erl_syntax_lib:analyze_module_attribute(ModuleAttribute) of
        Module when is_atom(Module) ->
            Module;
        {Module, _Params} when is_atom(Module) ->
            Module
    end.

get_file(Ast) ->
    case lists:keyfind(file, 3, Ast) of
        false ->
            undefined;
        FileAttribute ->
            {[_ | _] = File, _Line} = erl_syntax_lib:analyze_file_attribute(FileAttribute),
            File
    end.

%%-------------------------------------------------------------------
%% Internal Function Definitions - Tree Walking
%%-------------------------------------------------------------------

mapfold_ast_statement({function, Anno, Name, Arity, Clauses}, Context) ->
    #module_context{name = Module, file = File, sinks = Sinks} = Context,
    FunctionContext = #function_context{
        module = Module,
        file = File,
        name = Name,
        arity = Arity,
        sinks = Sinks
    },
    MappedClauses = [walk_function_statements(Clause, FunctionContext) || Clause <- Clauses],
    MappedStatement = {function, Anno, Name, Arity, MappedClauses},
    {MappedStatement, Context};
mapfold_ast_statement({attribute, _, record, {Name, Fields}} = Statement, Context) ->
    #module_context{pr_context = PrContext} = Context,
    FieldNames = record_field_names(Fields),
    UpdatedPrContext = fake_lager_pr:save_record_def(Name, FieldNames, PrContext),
    UpdatedContext = Context#module_context{pr_context = UpdatedPrContext},
    {Statement, UpdatedContext};
mapfold_ast_statement(Statement, Context) ->
    {Statement, Context}.

walk_function_statements(
    {call, Anno,
        {remote, RemoteAnno, {atom, _ModuleAnno, Module}, {atom, _FunctionAnno, Function}} =
            InvocationClause,
        Args},
    Context
) ->
    Arity = length(Args),
    MappedArgs = [walk_function_statements(Arg, Context) || Arg <- Args],
    SinkModules = Context#function_context.sinks,
    case
        lists:member(Module, SinkModules) andalso
            lists:member(Arity, [1, 2, 3]) andalso
            (lists:member(Function, ?LEVELS) orelse
                {true_but_unsafe, lists:keyfind(Function, 1, ?LEVELS_UNSAFE)})
    of
        true when
            % replace with `ok'
            Function =:= none
        ->
            {atom, RemoteAnno, ok};
        true ->
            transform_call(Anno, Function, MappedArgs, Module, Context);
        {true_but_unsafe, {_, Level}} ->
            transform_call(Anno, Level, MappedArgs, Module, Context);
        _ ->
            {call, Anno, InvocationClause, MappedArgs}
    end;
walk_function_statements(Node, Context) when is_tuple(Node) ->
    erl_syntax:revert(
        erl_syntax_lib:map_subtrees(
            fun(Child) -> walk_function_statements(Child, Context) end,
            Node
        )
    );
walk_function_statements(Node, _Context) ->
    Node.

%%-------------------------------------------------------------------
%% Internal Function Definitions - Transformation of Logging Calls
%%-------------------------------------------------------------------

transform_call(Anno, Level, Args, Module, Context) ->
    {Format, FormatArgs} = format_and_args(Anno, Args),
    Location = location_metadata(Anno, Context),
    CallMetadata = call_metadata(Anno, Args, Module),
    % Gate the whole call on a compile-time-inlined `logger:allow/2' check, exactly as
    % logger's own `?LOG_*' macros do (and as lager's original inlined level check did).
    % The format string, arguments and metadata textually live only in the `true' branch,
    % so they are evaluated solely when the level is enabled - no fun-wrapping or cost
    % heuristic needed, and side effects in the arguments fire exactly when the message is
    % emitted.
    %
    % The enabled branch dispatches through `logger:macro_log/5' rather than
    % `logger:Level/3': like the macros, this goes straight to `logger:log_allowed/4' and
    % skips the redundant `logger_config:allow/2' check that `logger:Level/3' (via
    % `do_log/3') would otherwise repeat. `macro_log/5' has been exported for exactly this
    % purpose since logger landed in OTP 21.
    %
    % The first argument is a populated `logger:location()' map carrying `mfa'/`file'/`line',
    % built exactly as logger's own `?LOCATION' macro does. Passing a populated location
    % (rather than `#{}') keeps the call within `macro_log/5's published contract -
    % `location()' mandates those three keys - so dialyzer/eqwalizer runs over downstream
    % transformed modules stay clean. Per logger's documented merge order the location is
    % also the lowest-priority metadata source, overridden by process and call-site metadata,
    % matching standard logger semantics.
    AllowCall =
        {call, Anno, {remote, Anno, {atom, Anno, logger}, {atom, Anno, allow}}, [
            {atom, Anno, Level},
            {atom, Anno, Context#function_context.module}
        ]},
    LogCall =
        {call, Anno, {remote, Anno, {atom, Anno, logger}, {atom, Anno, macro_log}}, [
            Location,
            {atom, Anno, Level},
            Format,
            FormatArgs,
            CallMetadata
        ]},
    {'case', Anno, AllowCall, [
        {clause, Anno, [{atom, Anno, true}], [], [LogCall]},
        {clause, Anno, [{atom, Anno, false}], [], [{atom, Anno, ok}]}
    ]}.

format_and_args(Anno, Args) ->
    case Args of
        [Format] ->
            % empty argument list
            {Format, {nil, Anno}};
        [Format, FormatArgs] ->
            {Format, FormatArgs};
        [_Metadata, Format, FormatArgs] ->
            {Format, FormatArgs}
    end.

%%-------------------------------------------------------------------
%% Internal Function Definitions - Logging Call Metadata
%%-------------------------------------------------------------------

% Builds the `logger:location()' map (`macro_log/5's first argument): the
% `mfa'/`file'/`line' triple, exactly as logger's `?LOCATION' macro. All three
% keys are mandatory in `location()', so `file' is always emitted - falling back
% to the empty string in the (practically impossible) absence of a `-file'
% attribute - to keep the literal a valid `location()' for downstream dialyzer.
location_metadata(Anno, Context) ->
    Line = erl_anno:line(Anno),
    File =
        case Context#function_context.file of
            undefined -> "";
            DefinedFile -> DefinedFile
        end,
    {map, Anno, [
        {map_field_assoc, Anno, {atom, Anno, mfa},
            {tuple, Anno, [
                {atom, Anno, Context#function_context.module},
                {atom, Anno, Context#function_context.name},
                {integer, Anno, Context#function_context.arity}
            ]}},
        {map_field_assoc, Anno, {atom, Anno, file}, {string, Anno, File}},
        {map_field_assoc, Anno, {atom, Anno, line}, {integer, Anno, Line}}
    ]}.

% Builds the call metadata (`macro_log/5's last argument): `lager_sink' for extra
% sinks plus any user-supplied metadata. The `mfa'/`file'/`line' triple lives in
% the location map instead (see location_metadata/2).
call_metadata(Anno, Args, Module) ->
    SinkMetadata = sink_metadata(Anno, Module),
    case Args of
        [_Fmt] ->
            SinkMetadata;
        [_Fmt, _FmtArgs] ->
            SinkMetadata;
        [ExtraMetadataList, _Fmt, _FmtArgs] ->
            extended_call_metadata(Anno, SinkMetadata, ExtraMetadataList)
    end.

sink_metadata(Anno, Module) ->
    Associations =
        % lager_sink => atom()
        case Module of
            lager ->
                [];
            CustomSink ->
                [{map_field_assoc, Anno, {atom, Anno, lager_sink}, {atom, Anno, CustomSink}}]
        end,
    {map, Anno, Associations}.

extended_call_metadata(Anno, BaseMetadata, ExtraMetadataList) ->
    case transform_metadata_list_into_map_field_associations(ExtraMetadataList) of
        {true, ExtraMetadataAssociations} ->
            {map, _, BaseMetadataAssociations} = BaseMetadata,
            {map, Anno, BaseMetadataAssociations ++ ExtraMetadataAssociations};
        false ->
            runtime_merged_extended_call_metadata(Anno, BaseMetadata, ExtraMetadataList)
    end.

transform_metadata_list_into_map_field_associations(ExtraMetadataList) ->
    transform_metadata_list_into_map_field_associations_recur(ExtraMetadataList, []).

transform_metadata_list_into_map_field_associations_recur(Clause, Acc) ->
    case Clause of
        {cons, ConsAnno, {tuple, _, [{atom, _, _} = KeyTerm, ValueTerm]}, NextClause} ->
            Assoc = {map_field_assoc, ConsAnno, KeyTerm, ValueTerm},
            UpdatedAcc = [Assoc | Acc],
            transform_metadata_list_into_map_field_associations_recur(NextClause, UpdatedAcc);
        {nil, _} ->
            Associations = lists:reverse(Acc),
            {true, Associations};
        _ ->
            false
    end.

runtime_merged_extended_call_metadata(Anno, BaseMetadata, ExtraMetadataList) ->
    % maps:merge(ExtraMetadata, BaseMetadata)
    {call, Anno, {remote, Anno, {atom, Anno, maps}, {atom, Anno, merge}},

        % Map1 - maps:from_list(ExtraMetadataList)
        [
            {call, Anno, {remote, Anno, {atom, Anno, maps}, {atom, Anno, from_list}}, [
                ExtraMetadataList
            ]},

            % Map2 - BaseMetadata
            BaseMetadata
        ]}.

%%-------------------------------------------------------------------
%% Internal Function Definitions - Pretty Printing of Records
%%-------------------------------------------------------------------

record_field_names(FieldsFromAst) ->
    lists:map(fun record_field_name/1, FieldsFromAst).

%% `record_field_name/1' copied from original `lager_transform',
%% licensed under Apache 2.0
record_field_name({record_field, _, {atom, _, FieldName}}) ->
    FieldName;
record_field_name({record_field, _, {atom, _, FieldName}, _Default}) ->
    FieldName;
record_field_name({typed_record_field, Field, _Type}) ->
    record_field_name(Field).

insert_pr_context_attribute([Attribute | Next], RecordDefs) ->
    case Attribute of
        {attribute, Line, module, _} = ModuleAttribute ->
            [
                ModuleAttribute,
                {attribute, Line, ?pr_context_attribute_name, [RecordDefs]}
                | Next
            ];
        OtherAttribute ->
            [
                OtherAttribute
                | insert_pr_context_attribute(Next, RecordDefs)
            ]
    end;
insert_pr_context_attribute([], _RecordDefs) ->
    % There's no module attribute - let the compiler handle it rather than have us crash
    [].

%%-------------------------------------------------------------------
%% Internal Function Definitions - Utiliities
%%-------------------------------------------------------------------

check_for_unsupported_options(Options) ->
    lists:foreach(
        fun
            ({lager_print_records_flag, true}) ->
                ok;
            ({Key, Value}) when
                Key =:= lager_truncation_size;
                Key =:= lager_print_records_flag;
                Key =:= lager_function_transforms
            ->
                error_logger:error_msg("[error] Unsupported option: '~s~n'", [{Key, Value}]),
                exit(normal);
            (_) ->
                ok
        end,
        Options
    ).

%write_terms(FilenameSuffix, List) ->
%    {attribute, _Anno, module, Module} = lists:keyfind(module, 3, List),
%    Filename = atom_to_list(Module) ++ "." ++ FilenameSuffix,
%    Format = fun(Term) -> io_lib:format("~tp.~n", [Term]) end,
%    Text = lists:map(Format, List),
%    file:write_file(Filename, Text).