-module(shake).
-export([start/1]).


% X_Y_[Q_]ZZ = SecretKey
%   Where
%     X  :: A (Alice) | B (Bob)
%     Y  :: P (Permanent) | E (Ephemeral)
%     Q  :: E (ec25519) | X (x25519)
%     ZZ :: ID | PK (Public Key) | S (Secret Key)

start(ContractID) ->
%   {A_P_ID, A_P_E_KP = #{public := A_P_E_PK, secret := A_P_E_SK}} = hz_key_master:make_key(<<>>),
    {_A_E_ID, A_E_E_KP = #{public := A_E_E_PK, secret := _A_E_E_SK}} = hz_key_master:make_key(<<>>),
    {B_P_ID,  B_P_E_KP = #{public := B_P_E_PK, secret := _B_P_E_SK}} = hz_key_master:make_key(<<>>),
    {_B_E_ID, B_E_E_KP = #{public := B_E_E_PK, secret := _B_E_E_SK}} = hz_key_master:make_key(<<>>),
    ServiceIDs = [B_P_ID],
    {_IPs, _Port, Salt, ServiceIDs} = resolve(ContractID, ServiceIDs),
    Salt = <<"Notional Salt">>,
    _A_String = connect_string(A_E_E_PK),
    _B_String = connect_string(B_P_E_PK, B_E_E_PK),
    % Key conversion
%   #{public := A_P_X_PK, secret := A_P_X_SK} = convert(A_P_E_KP),
    #{public := A_E_X_PK, secret := A_E_X_SK} = convert(A_E_E_KP),
    #{public := B_P_X_PK, secret := B_P_X_SK} = convert(B_P_E_KP),
    #{public := B_E_X_PK, secret := B_E_X_SK} = convert(B_E_E_KP),
    Version = proto_version(),
    A_InitSessionKey = session_key(a, A_E_X_SK, B_P_X_PK, B_E_X_PK, Version, Salt),
    B_InitSessionKey = session_key(b, B_P_X_SK, B_E_X_SK, A_E_X_PK, Version, Salt),
    A_InitSessionKey =:= B_InitSessionKey.


connect_string(PK1) ->
    <<(proto_name())/binary, ":", (proto_version())/binary, ":", PK1/binary>>.

connect_string(PK1, PK2) ->
    <<(proto_name())/binary, ":", (proto_version())/binary, ":", PK1/binary, PK2/binary>>.

proto_name() ->
    <<"GajuExpress">>.

proto_version() ->
    <<"001">>.


resolve(ContractID, IDs) ->
    ok = io:format("Would resolve ~p now...~n", [ContractID]),
    {["localhost", "127.0.0.1"], 7777, <<"Notional Salt">>, IDs}.


convert(#{public := PK, secret := <<Secret:32/binary, _/binary>>}) ->
    #{public => ed25519_pk_to_x25519(PK), secret => ed25519_sk_to_x25519(Secret)}.


%% Curve25519 Prime Field Constant: 2^255 - 19
%% Yes, in hex it read kind of like "lucky fed"
p() -> 16#7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFED.

ed25519_pk_to_x25519(<<ED25519_PK:32/binary>>) ->
    <<CompressedInt:256/little-integer>> = ED25519_PK,
    % Clear the sign bit (MSB) to get the raw y-coordinate
    Y = CompressedInt band ((1 bsl 255) - 1),
    
    % Compute u = (1 + y) / (1 - y) mod P
    Num = (1 + Y) rem p(),
    Den = (1 - Y + p()) rem p(),
    case Den =:= 0 of
        true -> 
            % If y == 1, the point maps to the point at infinity.
            % On X25519, this translates to u = 0.
            % TODO: Ask Peter whether this case should actually fail...?
            <<0:256/little-integer>>;
        false ->
            U = (Num * mod_inv(Den, p())) rem p(),
            <<U:256/little-integer>>
    end.

ed25519_sk_to_x25519(<<ED25519_SK_Secret:32/binary>>) ->
    %% Hash seed and slice the first 32 bytes (clamping done internally by crypto)
    <<X25519_SK:32/binary, _/binary>> = crypto:hash(sha512, ED25519_SK_Secret),
    X25519_SK.

mod_inv(A, M) ->
    {1, X, _} = ext_gcd(A, M),
    (X + M) rem M.

ext_gcd(A, 0) ->
    {A, 1, 0};
ext_gcd(A, B) ->
    {G, X1, Y1} = ext_gcd(B, A rem B),
    {G, Y1, X1 - (A div B) * Y1}.


session_key(a, A_E_X_SK, B_P_X_PK, B_E_X_PK, Version, Salt) ->
    DH_Static    = crypto:compute_key(ecdh, B_P_X_PK, A_E_X_SK, x25519),
    DH_Ephemeral = crypto:compute_key(ecdh, B_E_X_PK, A_E_X_SK, x25519),
    finalize_hkdf(DH_Static, DH_Ephemeral, Version, Salt);
session_key(b, B_P_X_SK, B_E_X_SK, A_E_X_PK, Version, Salt) ->
    DH_Static    = crypto:compute_key(ecdh, A_E_X_PK, B_P_X_SK, x25519),
    DH_Ephemeral = crypto:compute_key(ecdh, A_E_X_PK, B_E_X_SK, x25519),
    finalize_hkdf(DH_Static, DH_Ephemeral, Version, Salt).

finalize_hkdf(DH_Static, DH_Ephemeral, Version, Salt) ->
    MixedInput = <<DH_Static/binary, DH_Ephemeral/binary>>,
    Info       = <<"GajuExpress:", Version/binary, ":">>,
    hkdf(sha256, MixedInput, Salt, Info).


hkdf(Hash, IKM, Salt, Info) ->
    hkdf(Hash, IKM, Salt, Info, 32).

hkdf(Hash, IKM, Salt, Info, OutLen) ->
    PRK = extract(Hash, Salt, IKM),
    expand(Hash, PRK, Info, OutLen).

extract(Hash, <<>>, IKM) ->
    %% If salt is empty RFC 5869 dictates using a string of zeros equal to hash size
    Salt = binary:copy(<<0>>, hash_size(Hash)),
    extract(Hash, Salt, IKM);
extract(Hash, Salt, IKM) ->
    crypto:mac(hmac, Hash, Salt, IKM).

expand(Hash, PRK, Info, OutLen) ->
    HashLen = hash_size(Hash),
    BlockCount = (OutLen + HashLen - 1) div HashLen,
    case BlockCount > 255 of
        true ->
            error(hkdf_length_too_long);
        false -> 
            FullBlocks = expand_loop(Hash, PRK, Info, BlockCount, 1, <<>>, <<>>),
            %% Slice off excess trailing bytes to match exactly specified OutLen
            <<Output:OutLen/binary, _/binary>> = FullBlocks,
            Output
    end.

expand_loop(Hash, PRK, Info, N, Counter, PrevT, Acc) when Counter =< N ->
    %% Form block payload: T(i) = HMAC-Hash(PRK, T(i-1) | info | i)
    Payload = <<PrevT/binary, Info/binary, Counter:8>>,
    T = crypto:mac(hmac, Hash, PRK, Payload),
    expand_loop(Hash, PRK, Info, N, Counter + 1, T, <<Acc/binary, T/binary>>);
expand_loop(_, _, _, _, _, _, Acc) ->
    Acc.

hash_size(md5)    -> 16;
hash_size(sha)    -> 20;
hash_size(sha224) -> 28;
hash_size(sha256) -> 32;
hash_size(sha384) -> 48;
hash_size(sha512) -> 64.

    

% alice_start(ContractID) ->
%     {Pub, Sec} = crypto:generate_key(ecdh, secp256k1),
%     alice_start(Sec, Pub, ContractID).
% 
% alice_start(Sec, Pub, ContractID) ->
%     case net_lookup(ContractID) of
%         {ok, BobsIDs, Salt, Address, Port} -> connect(Sec, Pub, BobsIDs, Salt, Address, Port);
%         Error                              -> Error
%     end.
% 
% connect(Sec, Pub, BobsIDs, Salt, Address, Port) ->
%     Options = [{mode, binary}, {active, once}, {packet, 4}, {keepalive, true}],
%     case gen_tcp:connect(Address, Port, Options, 5000) of
%         {ok, Sock} -> push(Sec, Pub, BobsIDs, Salt, Sock);
%         Error      -> Error
%     end.
% 
% push(Sec, Pub, BobsIDs, Salt, Sock) ->
%     case gen_tcp:send(Sock, <<"BOBSandVAGEN:001:", Pub/binary>>) of
%         ok    -> await(Sec, BobsIDs, Salt, Sock);
%         Error -> Error
%     end.
% 
% await(Sec, BobsIDs, Salt, Sock) ->
%     ok = inet:setopts(Sock, [{active, once}]),
%     receive
%         {tcp, Sock, <<"BOBSandVAGEN:001:", B_S_Pub:65/binary, B_E_Pub:65/binary>>} ->
%             check(Sec, BobsIDs, B_S_Pub, B_E_Pub, Salt, Sock);
%         Other ->
%             ok = log(info, "Weird: ~p", [Other]),
%             {error, junk}
%         after 5000 ->
%             {error, timeout}
%     end.
% 
% check(A_E_Sec, BobsIDs, B_S_Pub, B_E_Pub, Salt, Sock) ->
%     case lists:member(B_S_Pub, BobsIDs) of
%         true  -> gen_session(A_E_Sec, B_S_Pub, B_E_Pub, Salt, Sock);
%         false -> {error, bad_id}
%     end.
% 
% gen_session(A_E_Sec, B_S_Pub, B_E_Pub, Salt, Sock) ->
%     case session_key(A_E_Sec, B_S_Pub, B_E_Pub, <<"001">>, Salt) of
%         {ok, SessionKey} -> loop(SessionKey, Sock);
%         Error            -> Error
%     end.