#include <cstdio>
#include <string>
#include <iostream>
#include <sstream>
#include <list>
#include <zmq.hpp>

/* TCPListener: declaration */
class TCPListener
{
public:
    TCPListener(int listen_port) : context(NULL), listenSocket(NULL), port(listen_port) {}
    virtual ~TCPListener();

    friend class MultiServer;

protected:
    // Methods called by MultiServer
    virtual std::string receive();
    virtual void setup(zmq::context_t *ctx);
    virtual void * getSocket() { return (void*)listenSocket; }

    zmq::context_t *context;
    zmq::socket_t *listenSocket;
    int port;
};

/* TCPListener: implementation */
TCPListener::~TCPListener()
{
    try
    {
        delete listenSocket;
        delete context;
    }
    catch (zmq::error_t &err)
    {
        printf("[!] TCPListener: %s\n", err.what());
    }
}

void TCPListener::setup(zmq::context_t *ctx)
{
    context = ctx;

    try
    {
        // Create and bind the listening socket
        listenSocket = new zmq::socket_t(*context, ZMQ_REP);

        std::stringstream bind_address;
        bind_address << "tcp://*:";
        bind_address << port;

        printf("[TCP] Binding to %s\n", bind_address.str().c_str());
        listenSocket->bind(bind_address.str().c_str());
    }
    catch (zmq::error_t &err)
    {
        printf("[!] TCPListener: %s\n", err.what());
        throw std::bad_alloc();
    }
}

std::string TCPListener::receive()
{
    zmq::message_t request;

    printf("[TCP] Receiving...\n");

    try
    {
        listenSocket->recv(&request);
    }
    catch (zmq::error_t &e)
    {
        printf("[!] IPC: receive() : %s\n", e.what());
    }

    // Debug
    printf ("Received message: \"%s\"\n", (char *)request.data());

    return std::string((const char*)request.data());
}

/* MultiServer: declaration */
class MultiServer
{
public:
    MultiServer() : context(1) {}            //!< Initialise the zmq context
    virtual ~MultiServer() {}

    void plug(TCPListener *listener);        //!< Plug a TCPListener
    std::string receive(long timeout = -1);  //!< Poll all listeners and receive data

protected:
    bool remListener(TCPListener *listener);  //!< Remove a TCPListener
    zmq::context_t context;                   //!< ZeroMQ context
    std::list<TCPListener *> listeners;       //!< This is where TCPListeners are added
};

/* MultiServer: implementation */
bool MultiServer::remListener(TCPListener *listener)
{
    bool found = false;

    std::list<TCPListener *>::iterator it;
    for (it = listeners.begin(); !found && it != listeners.end(); ++it)
    {
        if (*it == listener)
        {
            listeners.erase(it);
            printf("[-] Removed Listener %04x\n", listener);
            found = true;
        }
    }

    return found;
}

void MultiServer::plug(TCPListener *listener)
{
    // TODO: Prevent multiple plug to the same listener.
    if (listener != NULL)
    {
        listener->setup(&context);
        listeners.push_back(listener);
        printf("[+] Added TCPListener %04x (socket: %04x)\n", listener, listener->getSocket());
    }
}

std::string MultiServer::receive(long timeout)
{
    printf("[ ] Polling %ld listener(s)...\n", listeners.size());

    // Make a new pollitem array
    zmq::pollitem_t poll_items[listeners.size()];
    unsigned int i = 0;
    std::list<TCPListener *>::iterator it;
    for (it = listeners.begin(); it != listeners.end(); ++it, ++i)
    {
        poll_items[i].socket = (*it)->getSocket();
        poll_items[i].fd = 0;
        poll_items[i].events = ZMQ_POLLIN;
        poll_items[i].revents = 0;
    }

    try
    {
        for (unsigned int i = 0; i < listeners.size(); ++i)
            printf("[%u] { %04x | %d | %u | %u }\n",
                   i, poll_items[i].socket, poll_items[i].fd, poll_items[i].events, poll_items[i].revents);

        /*********** SEGMENTATION FAULT ***********/
        // Poll listeners
        zmq::poll(poll_items, (int)listeners.size(), timeout);
        /******************************************/
    }
    catch (zmq::error_t &err)
    {
        printf("[!] MultiServer: %s\n", err.what());
        throw std::bad_alloc();
    }

    /* -- Disabled for the moment --
    // Examine polled listeners and collect data
    for (unsigned int i = 0; i < listeners.size(); ++i)
    {
        if (pollItems[i]->revents & ZMQ_POLLIN)
        {
            printf("[ ] Incoming message on listener %d\n", i);

            // Fetch the message
            return listeners[i]->receive();
        }
    }
    */
    return "..."; // Temporary dummy return to make the compiler happy
}


/* ---   M A I N   --- */
int main(int argc, char *argv[])
{
    try
    {
        TCPListener tcp_listener(25000);
        MultiServer server;

        server.plug(&tcp_listener);

        while (true)
        {
            std::string message = server.receive();
            printf("[R] %s\n", message.c_str());
        }
    }
    catch (std::exception &e)
    {
        printf("[!] Exception: %s\n", e.what());
    }

    return 0;
}