%%%===================================================================
%%% matrix_megolm.erl — Megolm inbound session (fixed)
%%%===================================================================
-module(matrix_megolm).

-export([
    init_inbound/1, decrypt/2, session_id/1, pickle/1, unpickle/1,
    hmac/2, pkcs7unpad/1, decode_varint/1
]).

-record(mgm, {
    counter = 0 :: non_neg_integer(),
    r0 :: binary(), r1 :: binary(), r2 :: binary(), r3 :: binary(),
    ed25519_pub = undefined :: binary() | undefined
}).

-define(HKDF_INFO,   <<"MEGOLM_KEYS">>).
-define(MSG_VERSION, 3).
-define(KEY_VERSION, 2).

%%%===================================================================
%%% PUBLIC API
%%%===================================================================

init_inbound(SessionKeyBin) ->
    parse_session_key(SessionKeyBin).

decrypt(Session, CiphertextBin) ->
    try
        %% Protobuf-encoded MegOlm message:
        %%   VERSION(1) | 0x08 | INDEX(varint) | 0x12 | CT_LEN(varint) | CT | MAC(8) [| SIG(64)]
        %%
        %% 0x08 = protobuf field-1 tag (wire type 0 = varint)  → message index
        %% 0x12 = protobuf field-2 tag (wire type 2 = length-delimited) → ciphertext
        %%
        %% MAC covers every byte from the start of the message up to (but not including)
        %% the MAC itself; any trailing Ed25519 signature is NOT part of the MAC input.
        <<?MSG_VERSION:8, 16#08:8, R1/binary>> = CiphertextBin,
        {MsgIndex, <<16#12:8, R2/binary>>} = decode_varint(R1),
        {CtLen, R3} = decode_varint(R2),

        %% CtLen is the actual (post-padding) ciphertext length; always a multiple of 16.
        Ct = binary:part(R3, 0, CtLen),

        %% Compute exact MAC offset from the parsed structure — works regardless of
        %% whether a trailing Ed25519 signature is present.
        MacOffset = byte_size(CiphertextBin) - byte_size(R3) + CtLen,
        Mac       = binary:part(CiphertextBin, MacOffset, 8),
        MsgForMac = binary:part(CiphertextBin, 0, MacOffset),

        Session2 = advance(Session, MsgIndex),
        {AesKey, MacKey, AesIv} = derive_keys(Session2),

        ExpectedMac = binary:part(hmac(MacKey, MsgForMac), 0, 8),
        case ExpectedMac =:= Mac of
            false ->
                {error, mac_mismatch};
            true ->
                PlainPadded = crypto:crypto_one_time(aes_256_cbc, AesKey, AesIv, Ct, false),
                Plaintext   = pkcs7unpad(PlainPadded),
                Session3    = advance(Session2, MsgIndex + 1),
                {ok, {Plaintext, MsgIndex, Session3}}
        end
    catch C:R:_Stack ->
        {error, {decrypt_failed, C, R}}
    end.

%%%===================================================================
%%% Session helpers
%%%===================================================================

session_id(#mgm{ed25519_pub = undefined}) -> <<>>;
session_id(#mgm{ed25519_pub = Pub})       -> Pub.

pickle(S)    -> term_to_binary(S).
unpickle(Bin) ->
    try {ok, binary_to_term(Bin, [safe])}
    catch _:_ -> {error, bad_pickle}
    end.

%%%===================================================================
%%% Ratchet
%%%===================================================================

advance(S = #mgm{counter = Ctr}, Target) when Ctr >= Target -> S;
advance(S = #mgm{counter = Ctr, r0 = R0, r1 = R1, r2 = R2, r3 = R3}, Target) ->
    Next = Ctr + 1,
    {NR0, NR1, NR2, NR3} =
        if (Next band 16#FFFFFF) =:= 0 ->
               {hmac(R0, <<0>>), hmac(R0, <<1>>), hmac(R0, <<2>>), hmac(R0, <<3>>)};
           (Next band 16#FFFF) =:= 0 ->
               {R0, hmac(R1, <<1>>), hmac(R1, <<2>>), hmac(R1, <<3>>)};
           (Next band 16#FF) =:= 0 ->
               {R0, R1, hmac(R2, <<2>>), hmac(R2, <<3>>)};
           true ->
               {R0, R1, R2, hmac(R3, <<3>>)}
        end,
    advance(S#mgm{counter = Next, r0 = NR0, r1 = NR1, r2 = NR2, r3 = NR3}, Target).

%%%===================================================================
%%% Key derivation
%%%===================================================================

derive_keys(#mgm{r0 = R0, r1 = R1, r2 = R2, r3 = R3}) ->
    IKM  = <<R0/binary, R1/binary, R2/binary, R3/binary>>,
    Keys = hkdf(sha256, IKM, <<>>, ?HKDF_INFO, 80),
    <<AesKey:32/binary, MacKey:32/binary, AesIv:16/binary>> = Keys,
    {AesKey, MacKey, AesIv}.

%% HKDF-SHA256 (RFC 5869) — accumulates into a single binary, truncates to Len.
hkdf(Hash, IKM, Salt, Info, Len) ->
    PRK = crypto:mac(hmac, Hash, Salt, IKM),
    hkdf_expand(Hash, PRK, Info, 1, <<>>, <<>>, Len).

hkdf_expand(_Hash, _PRK, _Info, _I, _Prev, OKM, Len)
        when byte_size(OKM) >= Len ->
    binary:part(OKM, 0, Len);
hkdf_expand(Hash, PRK, Info, I, Prev, OKM, Len) ->
    T = crypto:mac(hmac, Hash, PRK, <<Prev/binary, Info/binary, I:8>>),
    hkdf_expand(Hash, PRK, Info, I + 1, T, <<OKM/binary, T/binary>>, Len).

%%%===================================================================
%%% Utilities
%%%===================================================================

hmac(Key, Data) -> crypto:mac(hmac, sha256, Key, Data).

pkcs7unpad(Bin) ->
    Len    = byte_size(Bin),
    PadLen = binary:last(Bin),
    case PadLen >= 1 andalso PadLen =< Len of
        true  -> binary:part(Bin, 0, Len - PadLen);
        false -> error({bad_padding, PadLen, Len})
    end.

decode_varint(Bin) -> decode_varint(Bin, 0, 0).

decode_varint(<<1:1, B:7, Rest/binary>>, Shift, Acc) ->
    decode_varint(Rest, Shift + 7, Acc bor (B bsl Shift));
decode_varint(<<0:1, B:7, Rest/binary>>, Shift, Acc) ->
    {Acc bor (B bsl Shift), Rest};
decode_varint(_, _, _) ->
    {0, <<>>}.

%%%===================================================================
%%% Session key parser
%%%===================================================================

parse_session_key(Bin) ->
    try
        case Bin of
            <<?KEY_VERSION:8, Counter:32/big-unsigned,
              R0:32/binary, R1:32/binary, R2:32/binary, R3:32/binary,
              Ed25519Pub:32/binary, _/binary>> ->
                {ok, #mgm{counter = Counter,
                          r0 = R0, r1 = R1, r2 = R2, r3 = R3,
                          ed25519_pub = Ed25519Pub}};
            <<?KEY_VERSION:8, Counter:32/big-unsigned,
              R0:32/binary, R1:32/binary, R2:32/binary, R3:32/binary, _/binary>> ->
                {ok, #mgm{counter = Counter,
                          r0 = R0, r1 = R1, r2 = R2, r3 = R3}}
        end
    catch _:_ ->
        {error, bad_session_key}
    end.