#!/usr/bin/python3
"""
@file
@ingroup hellorad
@copyright ESO - European Southern Observatory

@brief Simple client/server example.
"""

import sys
import argparse
import logging
import asyncio
import signal
import concurrent
from functools import partial

import zmq.asyncio
from rad.core.run import run_with_cleanup
from rad.services.msg import MsgRequestor
from rad.services.msg import MsgReplier

from rad.services.pubsub import Publisher
from rad.services.pubsub import Subscriber

from hellorad.pb_pb2 import *  # pylint: disable=E0611,W0401

logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)

def client(args, unknown):
    """
    Client sending request and subscribes to topics.
    """
    logger = logging.getLogger('client')

    parser = argparse.ArgumentParser()
    parser.add_argument('--req-endpoint', default='tcp://localhost:5555',
            dest='req_endpoint')
    parser.add_argument('--pub-endpoint', default='tcp://localhost:5556',
            dest='pub_endpoint')

    client_args = parser.parse_args(unknown)

    loop = asyncio.get_event_loop()
    zmq_ctx = zmq.asyncio.Context()

    requestor = MsgRequestor(endpoint = client_args.req_endpoint,
                             zmq_ctx = zmq_ctx,
                             loop = loop)

    queue = asyncio.Queue()

    def sub_handler(queue, topic_name, topic):
        logger.info('## Got topic: %s'%topic_name)
        queue.put_nowait((topic_name, topic))

    subscriber = Subscriber(handler=partial(sub_handler, queue),
                            topic_types=[ReqTest, RepTest],
                            zmq_ctx = zmq_ctx,
                            loop = loop)

    logger.info('Connecting to publisher %s', client_args.pub_endpoint)
    subscriber.connect(client_args.pub_endpoint)
    logger.info('Subscribing to "hellorad.hi" topic')
    subscriber.subscribe('hellorad.hi')
    # Start subscriber
    subtask = subscriber.start()

    async def send_request(loop, requestor):
        """
        Method to send a request.
        """
        logger.info('Sending request to %s', requestor)
        req = ReqTest()
        req.message = 'Ping pong'
        fut = requestor.send(req)
        # Wait for send future?
        rep = await asyncio.wait_for(requestor.receive(RepTest), timeout=3)
        logger.info('## Got reply: %s'%rep.message)

    async def stop_subscription(subscriber, queue):
        """
        Method that waits until a topic is received and then stops subscription.

        @todo: It would be nicer to have a coro-based subscriber instead of
        a callback-based one.
        """
        # Stop the subscription once we get a value
        topic = await queue.get()
        logger.info('Stopping subscriber')
        subscriber.stop()

    try:
        run_with_cleanup(send_request(loop, requestor),
                         subtask,
                         stop_subscription(subscriber, queue),
                         loop=loop, cleanup_timeout=5.0)
    except:
        logger.exception('Got exception')
    finally:
        logger.info('Exiting')

def server(args, unknown):
    """
    Server receiving requests and publishing topics.
    """
    logger = logging.getLogger('server')

    parser = argparse.ArgumentParser()
    parser.add_argument('--req-endpoint', default='tcp://*:5555',
            dest='req_endpoint')
    parser.add_argument('--sub-endpoint', default='tcp://*:5556',
            dest='pub_endpoint')

    server_args = parser.parse_args(unknown)

    loop = asyncio.get_event_loop()

    def sigint(loop, replier):
        """
        Handler for the SIGINT signal.
        """
        logger.info('SIGINT: Closing server')
        fut = replier.stop()

    def handler(replier, identity:bytes, msg):
        """
        Handler for requests.
        """
        logger.info('Got %s from %s'%(msg, identity.hex()))

        rep = RepTest()
        rep.message = 'Replier got: %s'%msg.message

        fut = replier.send(identity, rep)
        # Do I need to await the future to avoid warnings?
        # Give the client some time to receive the topics before shutting down
        loop.call_later(1.5, replier.stop)

    # MsgReplier
    logger.info('Binding request socket to %s'%server_args.req_endpoint)
    zmq_ctx = zmq.asyncio.Context()
    replier = MsgReplier(handler = handler,
                         endpoint = server_args.req_endpoint,
                         zmq_ctx = zmq_ctx,
                         msg_types = [ ReqTest ],
                         loop = loop)
    # Publisher
    logger.info('Binding pub socket to %s'%server_args.pub_endpoint)
    publisher = Publisher(endpoint = server_args.pub_endpoint,
                          zmq_ctx = zmq_ctx,
                          loop = loop)

    async def pub_stopper(replier, publisher):
        """
        Stop publishing if replier is shut down.
        """
        while not replier.is_closed():
            logger.info('Publishing topic "hellorad.hi"')
            msg = ReqTest()
            msg.message = 'Topic content'
            await publisher.publish('hellorad.hi', msg)
            logger.debug('Published ok. Sleeping for a bit')
            await asyncio.sleep(0.1)
            logger.debug('Back from sleep..')
        logger.info('Replier is done. Stopping publisher')
        publisher.stop()

    loop.add_signal_handler(signal.SIGINT, partial(sigint, loop, replier))

    try:
        run_with_cleanup(replier.start(),
                         pub_stopper(replier, publisher),
                         loop=loop, cleanup_timeout=3.0)
    except:
        logger.exception('Got exception')
    finally:
        logger.info('Exiting')

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Hello RAD for Python')
    parser.add_argument('type', choices=('server', 'client'))
    args, unknown = parser.parse_known_args()
    if args.type == 'server':
        server(args, unknown)

    if args.type == 'client':
        client(args, unknown)
    sys.exit(0)
