ex_stdlib

%%% @doc
%%% Base encoding and decoding module.
%%%
%%% This module provides functions for encoding and decoding data using
%%% various base encoding schemes: base16 (hex), base32, base64, and their
%%% URL-safe variants.
%%%
%%% Examples:
%%% ```
%%% Encoded = base:encode64("Hello, World!"),
%%% Decoded = base:decode64(Encoded),
%%% HexEncoded = base:encode16("binary data"),
%%% UrlSafe = base:url_encode64("data with + and /").
%%% '''
%%% @end
-module(base).

%% API - Base64
-export([encode64/1, decode64/1, decode64/2]).
-export([url_encode64/1, url_decode64/1, url_decode64/2]).

%% API - Base32
-export([encode32/1, decode32/1, decode32/2]).
-export([hex_encode32/1, hex_decode32/1, hex_decode32/2]).

%% API - Base16
-export([encode16/1, decode16/1, decode16/2]).

%% Types
-type encode_case() :: upper | lower.
-type padding() :: true | false.
-type decode_options() :: #{ignore => whitespace | none, letter_case => mixed | upper | lower}.

-export_type([encode_case/0, padding/0, decode_options/0]).

%%%=============================================================================
%%% Base64 Functions
%%%=============================================================================

%%% @doc
%%% Encodes binary data using Base64 encoding.
%%%
%%% Returns a binary containing the Base64-encoded data with padding.
%%% @end
-spec encode64(binary() | iodata()) -> binary().
encode64(Data) ->
    base64:encode(Data).

%%% @doc
%%% Decodes Base64-encoded data.
%%%
%%% Returns the decoded binary data. Raises an error if the input is invalid.
%%% @end
-spec decode64(binary() | string()) -> binary().
decode64(Data) ->
    try
        base64:decode(Data)
    catch
        error:Reason -> error({invalid_base64, Reason})
    end.

%%% @doc
%%% Decodes Base64-encoded data with options.
%%%
%%% Options:
%%% - ignore: whitespace | none - whether to ignore whitespace characters
%%% Returns {ok, Decoded} on success, or error on invalid input.
%%% @end
-spec decode64(binary() | string(), decode_options()) -> {ok, binary()} | error.
decode64(Data, Options) ->
    try
        CleanData = maybe_ignore_whitespace(Data, Options),
        Decoded = base64:decode(CleanData),
        {ok, Decoded}
    catch
        throw:has_whitespace -> error;
        _:_ -> error
    end.

%%% @doc
%%% Encodes binary data using URL-safe Base64 encoding.
%%%
%%% URL-safe Base64 uses '-' and '_' instead of '+' and '/', and omits padding.
%%% @end
-spec url_encode64(binary() | iodata()) -> binary().
url_encode64(Data) ->
    Encoded = base64:encode(Data),
    make_url_safe(Encoded).

%%% @doc
%%% Decodes URL-safe Base64-encoded data.
%%%
%%% Handles both padded and unpadded URL-safe Base64.
%%% @end
-spec url_decode64(binary() | string()) -> binary().
url_decode64(Data) ->
    try
        StandardData = make_standard_base64(Data),
        base64:decode(StandardData)
    catch
        error:Reason -> error({invalid_base64, Reason})
    end.

%%% @doc
%%% Decodes URL-safe Base64-encoded data with options.
%%%
%%% Returns {ok, Decoded} on success, or error on invalid input.
%%% @end
-spec url_decode64(binary() | string(), decode_options()) -> {ok, binary()} | error.
url_decode64(Data, Options) ->
    try
        CleanData = maybe_ignore_whitespace(Data, Options),
        StandardData = make_standard_base64(CleanData),
        Decoded = base64:decode(StandardData),
        {ok, Decoded}
    catch
        throw:has_whitespace -> error;
        _:_ -> error
    end.

%%%=============================================================================
%%% Base32 Functions
%%%=============================================================================

%%% @doc
%%% Encodes binary data using standard Base32 encoding.
%%%
%%% Uses the standard Base32 alphabet (A-Z, 2-7) with padding.
%%% @end
-spec encode32(binary() | iodata()) -> binary().
encode32(Data) ->
    encode_base32(iolist_to_binary(Data), standard_base32_alphabet()).

%%% @doc
%%% Decodes standard Base32-encoded data.
%%%
%%% Returns the decoded binary data. Raises an error if the input is invalid.
%%% @end
-spec decode32(binary() | string()) -> binary().
decode32(Data) ->
    try
        decode_base32(to_binary(Data), standard_base32_alphabet())
    catch
        error:Reason -> error({invalid_base32, Reason})
    end.

%%% @doc
%%% Decodes standard Base32-encoded data with options.
%%%
%%% Returns {ok, Decoded} on success, or error on invalid input.
%%% @end
-spec decode32(binary() | string(), decode_options()) -> {ok, binary()} | error.
decode32(Data, Options) ->
    try
        CleanData = maybe_ignore_whitespace(to_binary(Data), Options),
        Decoded = decode_base32(CleanData, standard_base32_alphabet()),
        {ok, Decoded}
    catch
        throw:has_whitespace -> error;
        _:_ -> error
    end.

%%% @doc
%%% Encodes binary data using Base32 hex encoding.
%%%
%%% Uses the extended hex alphabet (0-9, A-V) with padding.
%%% @end
-spec hex_encode32(binary() | iodata()) -> binary().
hex_encode32(Data) ->
    encode_base32(iolist_to_binary(Data), hex_base32_alphabet()).

%%% @doc
%%% Decodes Base32 hex-encoded data.
%%%
%%% Returns the decoded binary data. Raises an error if the input is invalid.
%%% @end
-spec hex_decode32(binary() | string()) -> binary().
hex_decode32(Data) ->
    try
        decode_base32(to_binary(Data), hex_base32_alphabet())
    catch
        error:Reason -> error({invalid_base32, Reason})
    end.

%%% @doc
%%% Decodes Base32 hex-encoded data with options.
%%%
%%% Returns {ok, Decoded} on success, or error on invalid input.
%%% @end
-spec hex_decode32(binary() | string(), decode_options()) -> {ok, binary()} | error.
hex_decode32(Data, Options) ->
    try
        CleanData = maybe_ignore_whitespace(to_binary(Data), Options),
        Decoded = decode_base32(CleanData, hex_base32_alphabet()),
        {ok, Decoded}
    catch
        throw:has_whitespace -> error;
        _:_ -> error
    end.

%%%=============================================================================
%%% Base16 Functions
%%%=============================================================================

%%% @doc
%%% Encodes binary data using Base16 (hexadecimal) encoding.
%%%
%%% Returns uppercase hexadecimal by default.
%%% @end
-spec encode16(binary() | iodata()) -> binary().
encode16(Data) ->
    encode_hex(iolist_to_binary(Data), upper).

%%% @doc
%%% Decodes Base16 (hexadecimal) encoded data.
%%%
%%% Accepts both uppercase and lowercase hexadecimal.
%%% @end
-spec decode16(binary() | string()) -> binary().
decode16(Data) ->
    try
        decode_hex(to_binary(Data))
    catch
        error:Reason -> error({invalid_base16, Reason})
    end.

%%% @doc
%%% Decodes Base16 (hexadecimal) encoded data with options.
%%%
%%% Options:
%%% - letter_case: mixed | upper | lower - expected case of hex digits
%%% Returns {ok, Decoded} on success, or error on invalid input.
%%% @end
-spec decode16(binary() | string(), decode_options()) -> {ok, binary()} | error.
decode16(Data, Options) ->
    try
        CleanData = maybe_ignore_whitespace(to_binary(Data), Options),
        validate_hex_case(CleanData, Options),
        Decoded = decode_hex(CleanData),
        {ok, Decoded}
    catch
        throw:has_whitespace -> error;
        _:_ -> error
    end.

%%%=============================================================================
%%% Internal functions - Base64
%%%=============================================================================

%% @private
make_url_safe(Base64) ->
    NoPadding = binary:replace(Base64, <<"=">>, <<>>, [global]),
    WithDash = binary:replace(NoPadding, <<"+">>, <<"-">>, [global]),
    binary:replace(WithDash, <<"/">>, <<"_">>, [global]).

%% @private
make_standard_base64(UrlSafeBase64) ->
    WithPlus = binary:replace(UrlSafeBase64, <<"-">>, <<"+">>, [global]),
    WithSlash = binary:replace(WithPlus, <<"_">>, <<"/">>, [global]),
    add_base64_padding(WithSlash).

%% @private
add_base64_padding(Data) ->
    case byte_size(Data) rem 4 of
        0 -> Data;
        2 -> <<Data/binary, "==">>;
        3 -> <<Data/binary, "=">>;
        _ -> error(invalid_base64_length)
    end.

%%%=============================================================================
%%% Internal functions - Base32
%%%=============================================================================

%% @private
standard_base32_alphabet() ->
    <<"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567">>.

%% @private
hex_base32_alphabet() ->
    <<"0123456789ABCDEFGHIJKLMNOPQRSTUV">>.

%% @private
encode_base32(<<>>, _Alphabet) ->
    <<>>;
encode_base32(Data, Alphabet) ->
    encode_base32_chunks(Data, Alphabet, <<>>).

%% @private
encode_base32_chunks(<<>>, _Alphabet, Acc) ->
    Acc;
encode_base32_chunks(Data, Alphabet, Acc) when byte_size(Data) < 5 ->
    % Pad the last chunk
    PaddedData = pad_to_5_bytes(Data),
    Chunk = encode_base32_chunk(PaddedData, Alphabet),
    PaddingSize = (5 - byte_size(Data)) * 8 div 5,
    TrimmedChunk = binary:part(Chunk, 0, 8 - PaddingSize),
    Padding = binary:copy(<<"=">>, PaddingSize),
    <<Acc/binary, TrimmedChunk/binary, Padding/binary>>;
encode_base32_chunks(<<Chunk:5/binary, Rest/binary>>, Alphabet, Acc) ->
    EncodedChunk = encode_base32_chunk(Chunk, Alphabet),
    encode_base32_chunks(Rest, Alphabet, <<Acc/binary, EncodedChunk/binary>>).

%% @private
encode_base32_chunk(<<A:8, B:8, C:8, D:8, E:8>>, Alphabet) ->
    % Convert 5 bytes (40 bits) to 8 base32 characters
    <<I1:5, I2:5, I3:5, I4:5, I5:5, I6:5, I7:5, I8:5>> = <<A:8, B:8, C:8, D:8, E:8>>,
    <<(binary:at(Alphabet, I1)), (binary:at(Alphabet, I2)),
      (binary:at(Alphabet, I3)), (binary:at(Alphabet, I4)),
      (binary:at(Alphabet, I5)), (binary:at(Alphabet, I6)),
      (binary:at(Alphabet, I7)), (binary:at(Alphabet, I8))>>.

%% @private
pad_to_5_bytes(Data) ->
    PaddingSize = 5 - byte_size(Data),
    <<Data/binary, 0:(PaddingSize * 8)>>.

%% @private
decode_base32(Data, Alphabet) ->
    CleanData = remove_padding(Data),
    decode_base32_chunks(CleanData, Alphabet, <<>>).

%% @private
decode_base32_chunks(<<>>, _Alphabet, Acc) ->
    Acc;
decode_base32_chunks(Data, Alphabet, Acc) when byte_size(Data) < 8 ->
    % Handle partial chunk
    PaddedData = <<Data/binary, (binary:copy(<<"A">>, 8 - byte_size(Data)))/binary>>,
    DecodedChunk = decode_base32_chunk(PaddedData, Alphabet),
    OutputSize = byte_size(Data) * 5 div 8,
    TrimmedChunk = binary:part(DecodedChunk, 0, OutputSize),
    <<Acc/binary, TrimmedChunk/binary>>;
decode_base32_chunks(<<Chunk:8/binary, Rest/binary>>, Alphabet, Acc) ->
    DecodedChunk = decode_base32_chunk(Chunk, Alphabet),
    decode_base32_chunks(Rest, Alphabet, <<Acc/binary, DecodedChunk/binary>>).

%% @private
decode_base32_chunk(<<C1, C2, C3, C4, C5, C6, C7, C8>>, Alphabet) ->
    I1 = base32_char_to_int(C1, Alphabet),
    I2 = base32_char_to_int(C2, Alphabet),
    I3 = base32_char_to_int(C3, Alphabet),
    I4 = base32_char_to_int(C4, Alphabet),
    I5 = base32_char_to_int(C5, Alphabet),
    I6 = base32_char_to_int(C6, Alphabet),
    I7 = base32_char_to_int(C7, Alphabet),
    I8 = base32_char_to_int(C8, Alphabet),

    <<A:8, B:8, C:8, D:8, E:8>> = <<I1:5, I2:5, I3:5, I4:5, I5:5, I6:5, I7:5, I8:5>>,
    <<A, B, C, D, E>>.

%% @private
base32_char_to_int(Char, Alphabet) ->
    case binary:match(Alphabet, <<Char>>) of
        {Index, 1} -> Index;
        nomatch -> error({invalid_base32_char, Char})
    end.

%%%=============================================================================
%%% Internal functions - Base16
%%%=============================================================================

%% @private
encode_hex(Data, upper) ->
    binary:encode_hex(Data).

%% @private
decode_hex(HexData) ->
    try
        binary:decode_hex(HexData)
    catch
        error:badarg -> error(invalid_hex)
    end.

%% @private
validate_hex_case(Data, Options) ->
    case maps:get(letter_case, Options, mixed) of
        mixed -> ok;
        upper -> validate_all_upper(Data);
        lower -> validate_all_lower(Data)
    end.

%% @private
validate_all_upper(<<>>) -> ok;
validate_all_upper(<<C, Rest/binary>>) when C >= $A, C =< $F; C >= $0, C =< $9 ->
    validate_all_upper(Rest);
validate_all_upper(_) ->
    error(invalid_case).

%% @private
validate_all_lower(<<>>) -> ok;
validate_all_lower(<<C, Rest/binary>>) when C >= $a, C =< $f; C >= $0, C =< $9 ->
    validate_all_lower(Rest);
validate_all_lower(_) ->
    error(invalid_case).

%%%=============================================================================
%%% Internal functions - Common
%%%=============================================================================

%% @private
to_binary(Data) when is_binary(Data) -> Data;
to_binary(Data) when is_list(Data) -> list_to_binary(Data).

%% @private
maybe_ignore_whitespace(Data, Options) ->
    case maps:get(ignore, Options, none) of
        whitespace -> remove_whitespace(Data);
        none ->
            case has_whitespace(Data) of
                true -> throw(has_whitespace);
                false -> Data
            end
    end.

%% @private
remove_whitespace(Data) ->
    binary:replace(Data, [<<" ">>, <<"\t">>, <<"\n">>, <<"\r">>], <<>>, [global]).

%% @private
has_whitespace(Data) ->
    binary:match(Data, [<<" ">>, <<"\t">>, <<"\n">>, <<"\r">>]) =/= nomatch.

%% @private
remove_padding(Data) ->
    binary:replace(Data, <<"=">>, <<>>, [global]).