-module(erltoken).

-on_load(init/0).

-export([
    list_encodings/0,
    encoding_for_model/1,
    count/2,
    encode/2,
    encode_with_special_tokens/2,
    decode/2,
    fits/3,
    trim_to_token_limit/3,
    remaining/3
]).

-export([
    list_encodings_nif/0,
    encoding_for_model_nif/1,
    count_nif/2,
    encode_nif/2,
    encode_with_special_tokens_nif/2,
    decode_nif/2
]).

-type name() :: binary().
-type token() :: non_neg_integer().

-spec init() -> ok | {error, term()}.
init() ->
    SoName = "erltoken_nif",
    case code:priv_dir(erltoken) of
        PrivDir when is_list(PrivDir) ->
            erlang:load_nif(filename:join(PrivDir, SoName), 0);
        {error, bad_name} ->
            erlang:load_nif(filename:join("priv", SoName), 0)
    end.

-spec list_encodings() -> [binary()].
list_encodings() ->
    list_encodings_nif().

-spec encoding_for_model(binary()) -> {ok, binary()} | {error, unknown_model | invalid_utf8}.
encoding_for_model(Model) when is_binary(Model) ->
    encoding_for_model_nif(Model).

-spec count(name(), binary()) -> {ok, non_neg_integer()} | {error, term()}.
count(Name, Text) when is_binary(Name), is_binary(Text) ->
    count_nif(Name, Text).

-spec encode(name(), binary()) -> {ok, [token()]} | {error, term()}.
encode(Name, Text) when is_binary(Name), is_binary(Text) ->
    encode_nif(Name, Text).

-spec encode_with_special_tokens(name(), binary()) -> {ok, [token()]} | {error, term()}.
encode_with_special_tokens(Name, Text) when is_binary(Name), is_binary(Text) ->
    encode_with_special_tokens_nif(Name, Text).

-spec decode(name(), [token()]) -> {ok, binary()} | {error, term()}.
decode(Name, Tokens) when is_binary(Name), is_list(Tokens) ->
    decode_nif(Name, Tokens).

-spec fits(name(), binary(), non_neg_integer()) -> {ok, boolean()} | {error, term()}.
fits(Name, Text, MaxTokens) when is_binary(Name), is_binary(Text), is_integer(MaxTokens), MaxTokens >= 0 ->
    case count(Name, Text) of
        {ok, Count} -> {ok, Count =< MaxTokens};
        {error, _} = Error -> Error
    end.

-spec trim_to_token_limit(name(), binary(), non_neg_integer()) -> {ok, binary()} | {error, term()}.
trim_to_token_limit(Name, Text, MaxTokens) when is_binary(Name), is_binary(Text), is_integer(MaxTokens), MaxTokens >= 0 ->
    case encode(Name, Text) of
        {ok, Tokens} when length(Tokens) =< MaxTokens ->
            {ok, Text};
        {ok, Tokens} ->
            decode(Name, lists:sublist(Tokens, MaxTokens));
        {error, _} = Error ->
            Error
    end.

-spec remaining(name(), binary(), non_neg_integer()) -> {ok, non_neg_integer()} | {error, term()}.
remaining(Name, Text, ContextWindow) when is_binary(Name), is_binary(Text), is_integer(ContextWindow), ContextWindow >= 0 ->
    case count(Name, Text) of
        {ok, Count} when Count =< ContextWindow -> {ok, ContextWindow - Count};
        {ok, _Count} -> {ok, 0};
        {error, _} = Error -> Error
    end.

list_encodings_nif() ->
    erlang:nif_error(nif_not_loaded).

encoding_for_model_nif(_Model) ->
    erlang:nif_error(nif_not_loaded).

count_nif(_Name, _Text) ->
    erlang:nif_error(nif_not_loaded).

encode_nif(_Name, _Text) ->
    erlang:nif_error(nif_not_loaded).

encode_with_special_tokens_nif(_Name, _Text) ->
    erlang:nif_error(nif_not_loaded).

decode_nif(_Name, _Tokens) ->
    erlang:nif_error(nif_not_loaded).