%% Copyright 2019 Ipregistry (https://ipregistry.co).
%%
%% Licensed under the Apache License, Version 2.0 (the "License");
%% you may not use this file except in compliance with the License.
%% You may obtain a copy of the License at
%%
%%     http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.

%% @doc Official Erlang client for the Ipregistry (https://ipregistry.co)
%% IP geolocation and threat data API.
%%
%% Create a client with {@link new/1} or {@link new/2}, then perform lookups:
%%
%% ```
%% Client = ipregistry:new(<<"YOUR_API_KEY">>),
%% {ok, Info} = ipregistry:lookup(Client, <<"8.8.8.8">>),
%% CountryName = ipregistry:get(Info, [location, country, name]).
%% '''
%%
%% Responses are returned as maps with binary keys, exactly as decoded from
%% the API's JSON. Use {@link get/2} and {@link get/3} to reach nested fields
%% conveniently.
%%
%% A client is a plain immutable map, so it is safe to share between
%% processes, embed in a supervisor's child state, or store in application
%% configuration.
-module(ipregistry).

-export([
    new/1, new/2,
    lookup/2, lookup/3,
    batch_lookup/2, batch_lookup/3,
    origin_lookup/1, origin_lookup/2,
    parse_user_agents/2,
    get/2, get/3,
    default_base_url/0,
    eu_base_url/0
]).

-ifdef(TEST).
-export([normalize_ip/1, build_query/1, cache_key/2, chunk/2, pmap/3]).
-endif.

-export_type([
    client/0,
    client_options/0,
    lookup_options/0,
    ip/0,
    ip_info/0,
    api_error/0,
    error_reason/0,
    batch_result/0
]).

-define(DEFAULT_BASE_URL, <<"https://api.ipregistry.co">>).
-define(EU_BASE_URL, <<"https://eu.api.ipregistry.co">>).
-define(MAX_BATCH_SIZE, 1024).

-type client() :: #{
    api_key := binary(),
    base_url := binary(),
    timeout := pos_integer(),
    max_retries := non_neg_integer(),
    retry_interval := pos_integer(),
    retry_on_server_error := boolean(),
    retry_on_too_many_requests := boolean(),
    max_batch_size := pos_integer(),
    batch_concurrency := pos_integer(),
    cache := atom() | pid() | undefined,
    user_agent := binary()
}.
%% An immutable client configuration created by {@link new/2}.

-type client_options() :: #{
    base_url => binary() | string(),
    timeout => pos_integer(),
    max_retries => non_neg_integer(),
    retry_interval => pos_integer(),
    retry_on_server_error => boolean(),
    retry_on_too_many_requests => boolean(),
    max_batch_size => pos_integer(),
    batch_concurrency => pos_integer(),
    cache => atom() | pid() | undefined,
    user_agent => binary() | string()
}.
%% Options accepted by {@link new/2}. See the README for the full reference.

-type lookup_options() :: #{
    fields => binary() | string(),
    hostname => boolean(),
    params => #{atom() | binary() | string() => binary() | string()}
}.
%% Per-request options: `fields' selects response fields using Ipregistry's
%% field selector syntax, `hostname' enables reverse-DNS resolution, and
%% `params' sets arbitrary extra query parameters.

-type ip() :: binary() | string() | inet:ip_address().
%% An IPv4 or IPv6 address, as a binary, a string, or an `inet' address tuple.

-type ip_info() :: #{binary() => term()}.
%% Decoded API response: a map with binary keys mirroring the JSON payload.

-type api_error() :: #{
    code := binary(),
    message := binary(),
    resolution := binary(),
    status => non_neg_integer()
}.
%% An error reported by the Ipregistry API. `code' is one of the codes listed
%% at https://ipregistry.co/docs/errors (for example `<<"INVALID_API_KEY">>').
%% `status' is the HTTP status, present when the whole request failed.

-type error_reason() ::
    {api_error, api_error()}
    | {client_error, term()}.
%% `api_error' means the API rejected the request; `client_error' covers
%% failures that happen on the client side (invalid input, transport errors,
%% undecodable responses).

-type batch_result() :: {ok, ip_info()} | {error, {api_error, api_error()}}.
%% One entry of a batch lookup; each entry succeeds or fails independently.

%%--------------------------------------------------------------------
%% Client construction
%%--------------------------------------------------------------------

%% @doc Creates a client with default settings.
%% @equiv new(ApiKey, #{})
-spec new(binary() | string()) -> client().
new(ApiKey) ->
    new(ApiKey, #{}).

%% @doc Creates a client authenticating with the given API key.
%%
%% You can obtain a key, along with a generous free tier, at
%% https://ipregistry.co. Ipregistry's OTP applications (`inets', `ssl') are
%% started if they are not running already.
%%
%% Raises `error({invalid_option, ...})' or `error(empty_api_key)' on
%% invalid configuration, since a misconfigured client is a programming
%% error rather than a runtime condition to handle.
-spec new(binary() | string(), client_options()) -> client().
new(ApiKey, Options) when is_map(Options) ->
    Key = to_binary(ApiKey),
    Key =:= <<>> andalso erlang:error(empty_api_key),
    {ok, _} = application:ensure_all_started(ipregistry),
    Defaults = #{
        base_url => ?DEFAULT_BASE_URL,
        timeout => 15000,
        max_retries => 3,
        retry_interval => 1000,
        retry_on_server_error => true,
        retry_on_too_many_requests => false,
        max_batch_size => ?MAX_BATCH_SIZE,
        batch_concurrency => 4,
        cache => undefined,
        user_agent => default_user_agent()
    },
    Client = maps:fold(fun apply_option/3, Defaults, Options),
    Client#{api_key => Key}.

%% @doc Returns the default API base URL, `https://api.ipregistry.co'.
-spec default_base_url() -> binary().
default_base_url() ->
    ?DEFAULT_BASE_URL.

%% @doc Returns the European Union API base URL,
%% `https://eu.api.ipregistry.co'. Pass it as the `base_url' option to
%% {@link new/2} to have your data processed in the EU only.
-spec eu_base_url() -> binary().
eu_base_url() ->
    ?EU_BASE_URL.

%%--------------------------------------------------------------------
%% Lookups
%%--------------------------------------------------------------------

%% @doc Looks up the data associated with the given IP address.
%% @equiv lookup(Client, Ip, #{})
-spec lookup(client(), ip()) -> {ok, ip_info()} | {error, error_reason()}.
lookup(Client, Ip) ->
    lookup(Client, Ip, #{}).

%% @doc Looks up the data associated with the given IP address.
%%
%% The address may be a binary, a string, or an `inet:ip_address()' tuple.
%% To look up the requester's own IP address use {@link origin_lookup/1}
%% instead. When the client is configured with a cache, a cache hit is
%% returned without contacting the API.
-spec lookup(client(), ip(), lookup_options()) -> {ok, ip_info()} | {error, error_reason()}.
lookup(Client, Ip, Options) ->
    case normalize_ip(Ip) of
        {ok, IpBin} ->
            Query = build_query(Options),
            Key = cache_key(IpBin, Query),
            case cache_get(Client, Key) of
                {ok, Info} ->
                    {ok, Info};
                miss ->
                    Path = with_query([<<"/">>, uri_string:quote(IpBin, ":")], Query),
                    case ipregistry_http:request(Client, get, Path, undefined) of
                        {ok, Info} ->
                            ok = cache_put(Client, Key, Info),
                            {ok, Info};
                        {error, _} = Error ->
                            Error
                    end
            end;
        {error, Reason} ->
            {error, {client_error, Reason}}
    end.

%% @doc Looks up several IP addresses at once.
%% @equiv batch_lookup(Client, Ips, #{})
-spec batch_lookup(client(), [ip()]) -> {ok, [batch_result()]} | {error, error_reason()}.
batch_lookup(Client, Ips) ->
    batch_lookup(Client, Ips, #{}).

%% @doc Looks up several IP addresses at once.
%%
%% Returns results in the same order as the input. Each entry may
%% independently succeed or fail (for example on an invalid address), so
%% inspect entries element by element:
%%
%% ```
%% {ok, Results} = ipregistry:batch_lookup(Client, [<<"8.8.8.8">>, <<"1.1.1.1">>]),
%% lists:foreach(
%%     fun({ok, Info}) -> io:format("~ts~n", [ipregistry:get(Info, [location, country, name])]);
%%        ({error, {api_error, #{code := Code}}}) -> io:format("failed: ~ts~n", [Code])
%%     end, Results).
%% '''
%%
%% The API accepts up to 1024 addresses per request; larger lists are
%% transparently split into several requests dispatched with bounded
%% concurrency (see the `max_batch_size' and `batch_concurrency' client
%% options) and results are reassembled in input order. An overall
%% `{error, ...}' return means the whole request failed (for example on
%% authentication or network errors), not that an individual entry did.
%%
%% Entries already present in the cache are served locally; only the
%% remainder are requested from the API.
-spec batch_lookup(client(), [ip()], lookup_options()) ->
    {ok, [batch_result()]} | {error, error_reason()}.
batch_lookup(Client, Ips, Options) when is_list(Ips) ->
    case normalize_ips(Ips, []) of
        {ok, IpBins} ->
            Query = build_query(Options),
            Entries = [
                {IpBin,
                    case cache_get(Client, cache_key(IpBin, Query)) of
                        {ok, Info} -> {cached, Info};
                        miss -> miss
                    end}
             || IpBin <- IpBins
            ],
            Misses = [IpBin || {IpBin, miss} <- Entries],
            case fetch_misses(Client, Misses, Query) of
                {ok, Fresh} ->
                    {ok, stitch(Client, Query, Entries, Fresh)};
                {error, _} = Error ->
                    Error
            end;
        {error, Reason} ->
            {error, {client_error, Reason}}
    end.

%% @doc Looks up the IP address the request originates from.
%% @equiv origin_lookup(Client, #{})
-spec origin_lookup(client()) -> {ok, ip_info()} | {error, error_reason()}.
origin_lookup(Client) ->
    origin_lookup(Client, #{}).

%% @doc Looks up the IP address the request originates from.
%%
%% The response additionally carries parsed User-Agent data under the
%% `<<"user_agent">>' key. Origin lookups are never cached, because the
%% requester IP is only known from the response.
-spec origin_lookup(client(), lookup_options()) -> {ok, ip_info()} | {error, error_reason()}.
origin_lookup(Client, Options) ->
    Query = build_query(Options),
    ipregistry_http:request(Client, get, with_query(<<"/">>, Query), undefined).

%% @doc Parses one or more raw User-Agent strings (such as the User-Agent
%% header of an incoming HTTP request) into structured data.
%%
%% Results preserve the order of the input, and each entry may independently
%% succeed or fail, like {@link batch_lookup/3} results.
-spec parse_user_agents(client(), [binary() | string()]) ->
    {ok, [batch_result()]} | {error, error_reason()}.
parse_user_agents(Client, UserAgents) when is_list(UserAgents) ->
    Body = iolist_to_binary(json:encode([to_binary(U) || U <- UserAgents])),
    case ipregistry_http:request(Client, post, <<"/user_agent">>, Body) of
        {ok, #{<<"results">> := Results}} when is_list(Results) ->
            {ok, [to_result_entry(R) || R <- Results]};
        {ok, Other} ->
            {error, {client_error, {unexpected_response, Other}}};
        {error, _} = Error ->
            Error
    end.

%%--------------------------------------------------------------------
%% Response access
%%--------------------------------------------------------------------

%% @doc Returns the value at the given path inside a response map, or
%% `undefined' when any path element is absent.
%% @equiv get(Info, Path, undefined)
-spec get(ip_info() | term(), [atom() | binary() | string()]) -> term().
get(Info, Path) ->
    get(Info, Path, undefined).

%% @doc Returns the value at the given path inside a response map, or
%% `Default' when any path element is absent.
%%
%% Path elements may be atoms, binaries, or strings; they are matched against
%% the response's binary keys:
%%
%% ```
%% ipregistry:get(Info, [location, country, name]).
%% ipregistry:get(Info, [security, is_vpn], false).
%% '''
-spec get(ip_info() | term(), [atom() | binary() | string()], term()) -> term().
get(Value, [], _Default) ->
    Value;
get(Map, [Key | Rest], Default) when is_map(Map) ->
    KeyBin = to_binary_key(Key),
    case Map of
        #{KeyBin := Value} -> get(Value, Rest, Default);
        _ -> Default
    end;
get(_Other, _Path, Default) ->
    Default.

%%--------------------------------------------------------------------
%% Client option handling
%%--------------------------------------------------------------------

apply_option(base_url, Value, Acc) ->
    Bin = to_binary(Value),
    Bin =:= <<>> andalso erlang:error({invalid_option, {base_url, Value}}),
    Acc#{base_url => strip_trailing_slash(Bin)};
apply_option(timeout, Value, Acc) when is_integer(Value), Value > 0 ->
    Acc#{timeout => Value};
apply_option(max_retries, Value, Acc) when is_integer(Value), Value >= 0 ->
    Acc#{max_retries => Value};
apply_option(retry_interval, Value, Acc) when is_integer(Value), Value > 0 ->
    Acc#{retry_interval => Value};
apply_option(retry_on_server_error, Value, Acc) when is_boolean(Value) ->
    Acc#{retry_on_server_error => Value};
apply_option(retry_on_too_many_requests, Value, Acc) when is_boolean(Value) ->
    Acc#{retry_on_too_many_requests => Value};
apply_option(max_batch_size, Value, Acc) when
    is_integer(Value), Value > 0, Value =< ?MAX_BATCH_SIZE
->
    Acc#{max_batch_size => Value};
apply_option(batch_concurrency, Value, Acc) when is_integer(Value), Value > 0 ->
    Acc#{batch_concurrency => Value};
apply_option(cache, Value, Acc) when is_atom(Value); is_pid(Value) ->
    Acc#{cache => Value};
apply_option(user_agent, Value, Acc) ->
    Bin = to_binary(Value),
    Bin =:= <<>> andalso erlang:error({invalid_option, {user_agent, Value}}),
    Acc#{user_agent => Bin};
apply_option(Key, Value, _Acc) ->
    erlang:error({invalid_option, {Key, Value}}).

strip_trailing_slash(Bin) ->
    case binary:last(Bin) of
        $/ -> strip_trailing_slash(binary:part(Bin, 0, byte_size(Bin) - 1));
        _ -> Bin
    end.

default_user_agent() ->
    Vsn =
        case application:get_key(ipregistry, vsn) of
            {ok, V} -> V;
            undefined -> "0.0.0"
        end,
    iolist_to_binary(["IpregistryClient/Erlang/", Vsn]).

%%--------------------------------------------------------------------
%% IP normalization
%%--------------------------------------------------------------------

normalize_ip(Ip) when is_binary(Ip) ->
    case Ip of
        <<>> -> {error, empty_ip};
        _ -> {ok, Ip}
    end;
normalize_ip(Ip) when is_list(Ip) ->
    case unicode:characters_to_binary(Ip) of
        Bin when is_binary(Bin) -> normalize_ip(Bin);
        _ -> {error, {invalid_ip, Ip}}
    end;
normalize_ip(Ip) when is_tuple(Ip) ->
    case catch inet:ntoa(Ip) of
        Str when is_list(Str) -> {ok, list_to_binary(Str)};
        _ -> {error, {invalid_ip, Ip}}
    end;
normalize_ip(Ip) ->
    {error, {invalid_ip, Ip}}.

normalize_ips([], Acc) ->
    {ok, lists:reverse(Acc)};
normalize_ips([Ip | Rest], Acc) ->
    case normalize_ip(Ip) of
        {ok, Bin} -> normalize_ips(Rest, [Bin | Acc]);
        {error, _} = Error -> Error
    end.

%%--------------------------------------------------------------------
%% Query parameters
%%--------------------------------------------------------------------

build_query(Options) when is_map(Options) ->
    Pairs = maps:fold(fun add_query_option/3, [], Options),
    case lists:keysort(1, Pairs) of
        [] -> <<>>;
        Sorted -> iolist_to_binary(uri_string:compose_query(Sorted))
    end.

add_query_option(fields, Value, Acc) ->
    [{<<"fields">>, to_binary(Value)} | Acc];
add_query_option(hostname, Value, Acc) when is_boolean(Value) ->
    [{<<"hostname">>, atom_to_binary(Value)} | Acc];
add_query_option(params, Params, Acc) when is_map(Params) ->
    maps:fold(
        fun(K, V, Inner) -> [{to_binary_key(K), to_binary(V)} | Inner] end,
        Acc,
        Params
    );
add_query_option(Key, Value, _Acc) ->
    erlang:error({invalid_option, {Key, Value}}).

with_query(Path, <<>>) ->
    iolist_to_binary(Path);
with_query(Path, Query) ->
    iolist_to_binary([Path, $?, Query]).

cache_key(IpBin, <<>>) ->
    IpBin;
cache_key(IpBin, Query) ->
    <<IpBin/binary, $;, Query/binary>>.

%%--------------------------------------------------------------------
%% Cache access
%%--------------------------------------------------------------------

cache_get(#{cache := undefined}, _Key) ->
    miss;
cache_get(#{cache := Cache}, Key) ->
    ipregistry_cache:get(Cache, Key).

cache_put(#{cache := undefined}, _Key, _Info) ->
    ok;
cache_put(#{cache := Cache}, Key, Info) ->
    ipregistry_cache:put(Cache, Key, Info).

%%--------------------------------------------------------------------
%% Batch machinery
%%--------------------------------------------------------------------

fetch_misses(_Client, [], _Query) ->
    {ok, []};
fetch_misses(Client, Misses, Query) ->
    #{max_batch_size := MaxBatchSize, batch_concurrency := Concurrency} = Client,
    Chunks = chunk(Misses, MaxBatchSize),
    ChunkResults = pmap(
        fun(Chunk) -> fetch_chunk(Client, Chunk, Query) end,
        Chunks,
        Concurrency
    ),
    case [Error || {error, _} = Error <- ChunkResults] of
        [FirstError | _] ->
            FirstError;
        [] ->
            {ok, lists:append([Entries || {ok, Entries} <- ChunkResults])}
    end.

fetch_chunk(Client, Chunk, Query) ->
    Body = iolist_to_binary(json:encode(Chunk)),
    Path = with_query(<<"/">>, Query),
    case ipregistry_http:request(Client, post, Path, Body) of
        {ok, #{<<"results">> := Results}} when is_list(Results) ->
            {ok, [to_result_entry(R) || R <- Results]};
        {ok, Other} ->
            {error, {client_error, {unexpected_response, Other}}};
        {error, _} = Error ->
            Error
    end.

%% Reassembles cached and freshly fetched entries in input order, caching
%% fresh successes along the way.
stitch(Client, Query, Entries, Fresh) ->
    stitch(Client, Query, Entries, Fresh, []).

stitch(_Client, _Query, [], _Fresh, Acc) ->
    lists:reverse(Acc);
stitch(Client, Query, [{_Ip, {cached, Info}} | Rest], Fresh, Acc) ->
    stitch(Client, Query, Rest, Fresh, [{ok, Info} | Acc]);
stitch(Client, Query, [{Ip, miss} | Rest], [Result | Fresh], Acc) ->
    case Result of
        {ok, Info} -> ok = cache_put(Client, cache_key(Ip, Query), Info);
        _ -> ok
    end,
    stitch(Client, Query, Rest, Fresh, [Result | Acc]);
stitch(Client, Query, [{_Ip, miss} | Rest], [], Acc) ->
    %% Defensive: the API returned fewer results than requested.
    Missing =
        {error,
            {api_error, #{
                code => <<>>,
                message => <<"missing result for requested IP address">>,
                resolution => <<>>
            }}},
    stitch(Client, Query, Rest, [], [Missing | Acc]).

to_result_entry(#{<<"code">> := Code} = Entry) when is_binary(Code) ->
    {error, {api_error, ipregistry_http:to_api_error(Entry)}};
to_result_entry(Entry) ->
    {ok, Entry}.

%% Runs Fun over Items with at most Limit concurrent workers, returning
%% results in input order.
pmap(Fun, Items, Limit) when is_integer(Limit), Limit >= 1 ->
    pmap_loop(Fun, lists:enumerate(Items), #{}, [], Limit).

pmap_loop(_Fun, [], Running, Acc, _Limit) when map_size(Running) =:= 0 ->
    [Result || {_Index, Result} <- lists:keysort(1, Acc)];
pmap_loop(Fun, [{Index, Item} | Pending], Running, Acc, Limit) when
    map_size(Running) < Limit
->
    Parent = self(),
    {Pid, Ref} = spawn_monitor(fun() -> Parent ! {pmap_result, self(), Fun(Item)} end),
    pmap_loop(Fun, Pending, Running#{Pid => {Ref, Index}}, Acc, Limit);
pmap_loop(Fun, Pending, Running, Acc, Limit) ->
    receive
        {pmap_result, Pid, Result} when is_map_key(Pid, Running) ->
            {{Ref, Index}, Running1} = maps:take(Pid, Running),
            erlang:demonitor(Ref, [flush]),
            pmap_loop(Fun, Pending, Running1, [{Index, Result} | Acc], Limit);
        {'DOWN', _Ref, process, Pid, Reason} when is_map_key(Pid, Running) ->
            {{_MonitorRef, Index}, Running1} = maps:take(Pid, Running),
            Result = {error, {client_error, {batch_worker_crashed, Reason}}},
            pmap_loop(Fun, Pending, Running1, [{Index, Result} | Acc], Limit)
    end.

chunk([], _Size) ->
    [];
chunk(Items, Size) when length(Items) =< Size ->
    [Items];
chunk(Items, Size) ->
    {Head, Tail} = lists:split(Size, Items),
    [Head | chunk(Tail, Size)].

%%--------------------------------------------------------------------
%% Small helpers
%%--------------------------------------------------------------------

to_binary(Value) when is_binary(Value) ->
    Value;
to_binary(Value) when is_list(Value) ->
    case unicode:characters_to_binary(Value) of
        Bin when is_binary(Bin) -> Bin;
        _ -> erlang:error({invalid_string, Value})
    end;
to_binary(Value) ->
    erlang:error({invalid_string, Value}).

to_binary_key(Key) when is_atom(Key) ->
    atom_to_binary(Key, utf8);
to_binary_key(Key) ->
    to_binary(Key).