#import <Foundation/Foundation.h>
#import <Network/Network.h>
#import <AppKit/AppKit.h>

#import <cassert>
#import <cstdlib>
#import <vector>

namespace nwdemo {

    // Dispatch queue our connections will be working on.
    dispatch_queue_t network_queue()
    {
        static dispatch_queue_t nwQueue = dispatch_queue_create("io.qt.nwdemo", DISPATCH_QUEUE_SERIAL);
        assert(nwQueue);
        return nwQueue;
    }

    // Keep it simple, use PSK (no certs).
    dispatch_data_t preshared_key()
    {
        const char bytes[] = "abcdefghijklmnopqrstuvwxyz";
        static dispatch_data_t psk = dispatch_data_create(bytes, sizeof bytes - 1, nullptr, DISPATCH_DATA_DESTRUCTOR_DEFAULT);
        assert(psk);
        return psk;
    }

    void add_psk(nw_protocol_options_t tlsOptions)
    {
        assert(tlsOptions);
        sec_protocol_options_t secOptions = nw_tls_copy_sec_protocol_options(tlsOptions);
        assert(secOptions);
        dispatch_data_t psk = preshared_key();
        assert(psk);
        sec_protocol_options_add_pre_shared_key(secOptions, psk, psk);
    }

    void add_transport_protocol(nw_parameters_t params)
    {
        assert(params);

        nw_protocol_options_t tcpOptions = nw_tcp_create_options();
        assert(tcpOptions);
        nw_protocol_stack_t stack = nw_parameters_copy_default_protocol_stack(params);
        assert(stack);
        nw_protocol_stack_set_transport_protocol(stack, tcpOptions);
    }

    // This framer inserts TLS protocol into the stack and then is marked as ready.
    const auto tlsStarter = ^nw_framer_start_result_t(nw_framer_t  _Nonnull framer)
    {
        NSLog(@"framer (TLS-inserter) %@ starting ...", framer);
        nw_protocol_options_t tlsOptions = nw_tls_create_options();
        add_psk(tlsOptions);
        nw_framer_prepend_application_protocol(framer, tlsOptions);
        
        nw_framer_pass_through_input(framer);
        nw_framer_pass_through_output(framer);
        
        // Emulate the delay with TLS handshake, pretend we write
        // 'start TLS' in 'clear text' before TLS handshake:
        dispatch_async(dispatch_get_main_queue(), ^{
            nw_framer_async(framer, ^{
                NSLog(@"marking framer %@ as 'ready'", framer);
                nw_framer_mark_ready(framer);
            });
        });
        return nw_framer_start_result_will_mark_ready;
    };

    // Add a custom framer below TLS.
    void insert_custom_protocol(nw_parameters_t params)
    {
        assert(params);
        nw_protocol_stack_t stack = nw_parameters_copy_default_protocol_stack(params);
        assert(stack);

        nw_protocol_stack_clear_application_protocols(stack);
        
        nw_protocol_definition_t framerDefinition = nw_framer_create_definition("tls.starter", 0, tlsStarter);
        assert(framerDefinition);
        nw_protocol_options_t customOptions = nw_framer_create_options(framerDefinition);
        assert(customOptions);
        nw_protocol_stack_prepend_application_protocol(stack, customOptions);
    }

    nw_parameters_t connection_parameters()
    {
        nw_parameters_t params = nw_parameters_create();
        assert(params);
        add_transport_protocol(params);
        insert_custom_protocol(params);
        return params;
    }

    // Our server and clients start as generic connections, with TCP as transport and custom protocol,
    // aka framer, which injects TLS to application protocol stack to enable encryption.
    nw_listener_t listener = nullptr;
    void start_server()
    {
        assert(!listener);

        dispatch_queue_t queue = network_queue();
        assert(queue);

        nw_parameters_t connectionParameters = connection_parameters();
        assert(connectionParameters);

        nw_endpoint_t localEndpoint = nw_endpoint_create_host("127.0.0.1", "22345");
        assert(localEndpoint);
        nw_parameters_set_local_endpoint(connectionParameters, localEndpoint);

        listener = nw_listener_create(connectionParameters);
        assert(listener);
        nw_listener_set_queue(listener, queue);

        nw_listener_set_state_changed_handler(listener, ^(nw_listener_state_t state, nw_error_t error) {
            NSLog(@"server: listener state changed on a queue %s", dispatch_queue_get_label(nullptr));
            switch (state) {
            case nw_listener_state_waiting:
                NSLog(@"server: listener is in waiting state ...");
                break;
            case nw_listener_state_failed:
                NSLog(@"server: listener failed %@ %d", error, nw_error_get_error_code(error));
                break;
            case nw_listener_state_ready:
                NSLog(@"server: listener is ready on port %u", nw_listener_get_port(listener));
                break;
            case nw_listener_state_cancelled:
                NSLog(@"server:listener was cancelled");
                break;
            default:;
            }
        });
   
        // Handler of incoming connections.
        nw_listener_set_new_connection_handler(listener, ^(nw_connection_t connection) {
            void process_incoming_connection(nw_connection_t connection);
            process_incoming_connection(connection);
        });

        // Listen for incoming connections.
        nw_listener_start(listener);
        NSLog(@"server: listening on port %d", nw_listener_get_port(listener));
    }

    std::vector<nw_connection_t> clients;
    void start_client()
    {
        dispatch_queue_t queue = network_queue();
        assert(queue);

        nw_parameters_t connectionParams = connection_parameters();
        assert(connectionParams);

        nw_endpoint_t endPoint = nw_endpoint_create_host("127.0.0.1", "22345");
        assert(endPoint);
        nw_connection_t clientConnection = nw_connection_create(endPoint, connectionParams);
        assert(clientConnection);
        nw_connection_set_queue(clientConnection, queue);

        // Connection state handler.
        auto stateHandler = ^(nw_connection_state_t state, nw_error_t  _Nullable error) {
            NSLog(@"client: connection state change on %s", dispatch_queue_get_label(nullptr));

            switch (state) {
            case nw_connection_state_failed:
            case nw_connection_state_invalid:
            case nw_connection_state_cancelled:
                NSLog(@"client: connection failed with error %@", error);
                break;
            case nw_connection_state_ready:
                // This means TLS handshake was successful.
                NSLog(@"client: connection is ready (encrypted) ...");
                break;
            default:
                NSLog(@"client: connection is preparing/waiting ...");
            }
        };

        nw_connection_set_state_changed_handler(clientConnection, stateHandler);

        // Connect.
        nw_connection_start(clientConnection);
        clients.push_back(clientConnection);
    }

    std::vector<nw_connection_t> incomingConnections;
    void process_incoming_connection(nw_connection_t connection)
    {
        assert(connection);
        incomingConnections.push_back(connection);
        
        dispatch_queue_t queue = network_queue();
        assert(queue);
        nw_connection_set_queue(connection, queue);

        auto stateHandler = ^(nw_connection_state_t state, nw_error_t  _Nullable error) {
            switch (state) {
            case nw_connection_state_failed:
            case nw_connection_state_invalid:
            case nw_connection_state_cancelled:
                NSLog(@"server: connection failed with error %@", error);
                break;
            case nw_connection_state_ready:
                // This means TLS handshake was successful.
                NSLog(@"server: connection %@ is ready.", connection);
                {
                    //nw_parameters_t params = nw_connection_copy_parameters(connection);
                    //nw_protocol_stack_t
                }
                break;
            default:
                NSLog(@"server: connection is preparing/waiting ...");
            }
        };

        nw_connection_set_state_changed_handler(connection, stateHandler);
        
        NSLog(@"server: processing incoming connection %@", connection);
        nw_connection_start(connection);
    }
}

int main(int argc, const char ** argv)
{
    using namespace nwdemo;
    
    start_server();
    start_client();
    start_client();

    return NSApplicationMain(argc, argv);
}
