#include "client.h"
#include "network/tcp-chat-protocol/v1.h"

void ServerConnection_close(ServerConnection* conn){
    if(conn == NULL)
        return;
    RSA_destroyPublicKey(&conn->server_pk);
    EncryptedSocketTCP_destroy(&conn->sock);
    free(conn->session_key.data);
    free(conn);
}

/// @brief 
/// @param server_link_cstr address:port:public_key
/// @return 
Result(void) ServerLink_parse(cstr server_link_cstr, EndpointIPv4* server_end_out, br_rsa_public_key* server_key_out){
    Deferral(8);
    str server_link_str = str_from_cstr(server_link_cstr);
    
    // parse address and port
    i32 sep_pos = str_seekChar(server_link_str, ':', 0);
    if(sep_pos == -1){
        Return RESULT_ERROR_FMT("server link is invalid: %s", server_link_cstr);
    }
    *server_end_out = EndpointIPv4_INVALID;
    try_void(EndpointIPv4_parse(server_link_cstr, server_end_out));    
    if(EndpointIPv4_is_invalid(*server_end_out)){
        Return RESULT_ERROR_FMT("server address or port is invalid: %s", server_link_cstr);
    }

    // parse public key
    sep_pos = str_seekChar(server_link_str, ':', sep_pos + 1);
    if(sep_pos == -1){
        Return RESULT_ERROR_FMT("server link is invalid: %s", server_link_cstr);
    }
    str server_key_str = str_sliceAfter(server_link_str, sep_pos + 1);
    char* server_key_cstr = str_copy(server_key_str).data;
    Defer(free(server_key_cstr));
    try_void(RSA_parsePublicKey_base64(server_key_cstr, server_key_out));
    
    Return RESULT_VOID;
}

Result(ServerConnection*) ServerConnection_open(ClientCredentials* client_credentials, cstr server_link_cstr){
    Deferral(64);

    ServerConnection* conn = (ServerConnection*)malloc(sizeof(ServerConnection));
    memset(conn, 0, sizeof(*conn));
    bool success = false;
    Defer(
        if(!success)
            ServerConnection_close(conn);
    );

    try_void(ServerLink_parse(server_link_cstr, &conn->server_end, &conn->server_pk));
    RSAEncryptor_construct(&conn->rsa_enc, &conn->server_pk);
    
    conn->session_key = Array_alloc_size(AES_SESSION_KEY_SIZE);
    // generate random session key
    br_hmac_drbg_context key_rng = { .vtable = &br_hmac_drbg_vtable };
    rng_init_sha256_seedFromSystem(&key_rng.vtable);
    br_hmac_drbg_generate(&key_rng, conn->session_key.data, conn->session_key.size);

    // connect to server address
    try(Socket _s, i, socket_open_TCP());
    // TODO: set socket timeout to 5 seconds
    try_void(socket_connect(_s, conn->server_end));
    EncryptedSocketTCP_construct(&conn->sock, _s, NETWORK_BUFFER_SIZE, conn->session_key);
    
    Array(u8) buffer = Array_alloc_size(NETWORK_BUFFER_SIZE);
    // fix for valgrind false detected errors about uninitialized memory
    Array_memset(buffer, 0xCC);
    Defer(free(buffer.data));

    // construct PacketHeader and ClientHandshake in buffer
    PacketHeader_construct(buffer.data, PROTOCOL_VERSION,
        PacketType_ClientHandshake, sizeof(ClientHandshake));
    ClientHandshake_construct(
        Array_sliceAfter(buffer, sizeof(PacketHeader)).data,
        conn->session_key);
    u32 header_and_message_size = sizeof(PacketHeader) + sizeof(ClientHandshake);
    // encrypt message by server public key
    Array(u8) bufferPart_encryptedClientHandshake = Array_sliceAfter(buffer, header_and_message_size);
    try(u32 rsa_enc_size, u, 
        RSAEncryptor_encrypt(
            &conn->rsa_enc,
            Array_sliceBefore(buffer, header_and_message_size),
            bufferPart_encryptedClientHandshake
        )
    );
    bufferPart_encryptedClientHandshake.size = rsa_enc_size;
    // send encrypted message
    try_void(socket_send(conn->sock.sock, bufferPart_encryptedClientHandshake));
    
    // receive server response
    try_void( 
        EncryptedSocketTCP_recv(&conn->sock,
            Array_sliceBefore(buffer, sizeof(PacketHeader)),
            SocketRecvFlag_WaitAll
        )
    );
    PacketHeader* packet_header = buffer.data;
    try_void(PacketHeader_validateMagic(packet_header));
    
    // handle server response
    switch(packet_header->type){
        case PacketType_ErrorMessage: {
            u32 err_msg_size = packet_header->content_size;
            if(err_msg_size > conn->sock.recv_buf.size)
                err_msg_size = conn->sock.recv_buf.size;
            Array(u8) err_buf = Array_alloc_size(err_msg_size + 1);
            bool err_msg_completed = false;
            Defer(
                if(!err_msg_completed)
                    free(err_buf.data);
            );

            // receive error message
            try_void(
                EncryptedSocketTCP_recv(
                    &conn->sock,
                    Array_sliceBefore(err_buf, err_msg_size),
                    SocketRecvFlag_WaitAll
                )
            );

            ((u8*)err_buf.data)[err_msg_size] = 0;
            err_msg_completed = true;
            Return RESULT_ERROR((char*)err_buf.data, true);
        }
        case PacketType_ServerHandshake: {
            Array(u8) bufferPart_ServerHandshake = { 
                .data = (u8*)buffer.data + sizeof(PacketHeader),
                .size = sizeof(ServerHandshake)
            };
            try_void(
                EncryptedSocketTCP_recv(
                    &conn->sock,
                    bufferPart_ServerHandshake,
                    SocketRecvFlag_WaitAll
                )
            );
            ServerHandshake* server_handshake = bufferPart_ServerHandshake.data;
            conn->session_id = server_handshake->session_id;
            break;
        }
        default:
            Return RESULT_ERROR_FMT("unexpected response type: %i", packet_header->type);
    }
    
    success = true;
    Return RESULT_VALUE(p, conn);
}