#
# Copyright 2018 3liz
# Author David Marteau
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

""" Zmq asynchrone client
"""

import asyncio
import logging
import pickle
import sys
import traceback
import uuid

from typing import (
    Dict,
    Mapping,
    Optional,
    Tuple,
    cast,
)

import zmq
import zmq.asyncio

from ..logger import setup_log_handler
from .messages import RequestMessage

LOGGER = logging.getLogger('SRVLOG')


class RequestTimeoutError(Exception):
    pass


class RequestGatewayError(Exception):
    pass


class RequestProxyError(Exception):
    pass


class AsyncResponseHandler:
    def __init__(self, correlation_id: bytes):

        loop = asyncio.get_running_loop()

        self.correlation_id = correlation_id
        self.headers: Dict[str, str] = {}
        # Create a future for sending the result
        self._future = loop.create_future()
        self._chunks: Optional[asyncio.Queue] = None
        self._has_more: bool = False
        self.data: Optional[bytes] = None
        self.metadata: Optional[Dict[str, str]] = None
        self.status: int = -1

    def _set_exception(self, exc: Exception):
        self._has_more = False
        self._future.set_exception(exc)

    def _set_result(self, data: bytes):
        """ Set raw results from request

            This method send the result to the stored
            future.
        """
        if self.status == -1:
            status, hdrs, body, meta = pickle.loads(data)
            self.headers = hdrs
            self.data = body
            self.status = status
            self.metadata = meta
            # We are waiting for more data
            # Create a queue to collect the remaining chunks
            if self.status == 206:
                self._has_more = True
                self._chunks = asyncio.Queue()
            # Send the result
            self._future.set_result(self)
        elif self._has_more:
            body, has_more, extra = pickle.loads(data)
            self._has_more = has_more
            chunks = cast(asyncio.Queue, self._chunks)
            chunks.put_nowait((body, has_more))
            self.extra = extra

    def _done(self) -> bool:
        """ Check if there is more data to come
        """
        return not self._has_more and self._future.done()

    async def _get(self, timeout: int) -> 'AsyncResponseHandler':
        """ wait for result and return the parsed
            result
        """
        try:
            return await asyncio.wait_for(self._future, timeout)
        except asyncio.TimeoutError:
            self._has_more = False
            raise RequestTimeoutError()

    async def _next_chunk(self, timeout: int) -> Tuple[bytes, bool]:
        """ Get next chunk
        """
        try:
            chunks = cast(asyncio.Queue, self._chunks)
            return await asyncio.wait_for(chunks.get(), timeout)
        except asyncio.TimeoutError:
            self._has_more = False
            raise RequestTimeoutError()


class AsyncClient:
    """ Async DEALER ZMQ client
    """

    def __init__(self, address: str, identity: Optional[bytes] = None):

        context = zmq.asyncio.Context.instance()

        self.identity = identity or uuid.uuid1().bytes

        sock = context.socket(zmq.DEALER)
        sock.setsockopt(zmq.LINGER, 500)    # Needed for socket no to wait on close
        sock.setsockopt(zmq.IMMEDIATE, 1)   # Do not queue if there is no connection
        sock.identity = self.identity
        sock.connect(address)

        self._running = False
        self._handlers: Dict[bytes, AsyncResponseHandler] = {}
        self._socket = sock
        self._polling = False
        self._poll_task: Optional[asyncio.Task] = None
        LOGGER.info("Starting client %s", self.identity)

    async def _poll(self):
        """ Handle incoming messages
        """
        self._polling = True
        while self._handlers:
            try:
                correlation_id, data, *rest = await self._socket.recv_multipart()
                # Get if there is a future pending for that message
                try:
                    handler = self._handlers[correlation_id]
                    if rest and data == b'ERR':
                        handler._set_exception(RequestProxyError(rest[0]))
                    else:
                        handler._set_result(data)
                    # Remove handlers from the heap if we are done
                    if handler._done():
                        self._handlers.pop(correlation_id, None)
                except KeyError:
                    LOGGER.warning("%s: No pending future found for message %s", self.identity, correlation_id)
            except zmq.ZMQError as err:
                LOGGER.error("%s error:  zmq error: %s (%s)", self.identity, zmq.strerror(err.errno), err.errno)
            except Exception as err:
                LOGGER.error("%s exception %s\n%s", self.identity, err, traceback.format_exc())
        self._polling = False
        self._poll_task = None

    async def fetch(
        self, query: str,
        method: str = 'GET',
        headers: Mapping[str, str] = {},
        data: Optional[bytes] = None,
        timeout: int = 5,
    ) -> AsyncResponseHandler:
        """ Send a request message to the worker
        """
        # Send request
        request = pickle.dumps(RequestMessage(query, headers=headers, method=method, data=data), -1)
        correlation_id = uuid.uuid1().bytes
        assert correlation_id not in self._handlers
        try:
            await self._socket.send_multipart([correlation_id, request], flags=zmq.DONTWAIT)
        except zmq.ZMQError as err:
            LOGGER.error("%s (%s)", zmq.strerror(err.errno), err.errno)
            raise RequestGatewayError()

        # Create response handler and register it
        handler = AsyncResponseHandler(correlation_id)
        self._handlers[correlation_id] = handler
        # Run poller if needed
        if not self._polling:
            self._poll_task = asyncio.create_task(self._poll())

        # Wait for response
        try:
            return await handler._get(timeout)
        except Exception:
            self._handlers.pop(correlation_id, None)
            raise

    async def fetch_more(self, response, timeout=5):
        """ Request next chunk
        """
        try:
            while True:
                data, has_more = await response._next_chunk(timeout)
                if not has_more:
                    break
                yield data
        except Exception:
            self._handlers.pop(response.correlation_id, None)
            raise

    def terminate(self):
        LOGGER.info("Terminating client %s", self.identity)
        self._running = False
        self._handlers = {}
        self._socket.close()


if __name__ == '__main__':
    import argparse
    import signal

    from time import sleep

    parser = argparse.ArgumentParser(description='Test Client')
    parser.add_argument('--host', metavar="host", default='tcp://localhost', help="router host")
    parser.add_argument('--router', metavar='address', default='{host}:8880', help="router address")
    parser.add_argument('--logging', choices=['debug', 'info', 'warning', 'error'], default='info',
                        help="set log level")
    parser.add_argument('--identity', default='', help="Set worker identity")
    parser.add_argument('--count', default=1, type=int, help="Number of requests")

    args = parser.parse_args()

    setup_log_handler(args.logging)
    print(f"Log level set to {logging.getLevelName(LOGGER.level)}\n", file=sys.stderr)  # noqa: T201

    LOGGER.setLevel(getattr(logging, args.logging.upper()))

    client = AsyncClient(args.router.format(host=args.host), bytes(args.identity.encode('ascii')))
    sleep(1)  # Give some time to connection to establish

    async def fetch(index):
        try:
            response = await client.fetch(query="?service=WMS", data=b"Hello world from %d" % index)
            print("%d -> response = %s" % (index, response.data))  # noqa: T201
            async for chunk in client.fetch_more(response):
                print("%d -> chunk = %s" % (index, chunk))  # noqa: T201
        except RequestTimeoutError:
            LOGGER.error("%d -> TIMEOUT", index)
        except RequestGatewayError:
            LOGGER.error("%d -> GATEWAY ERROR", index)

    loop = asyncio.new_event_loop()
    loop.add_signal_handler(signal.SIGINT, loop.stop)
    loop.run_until_complete(asyncio.wait([fetch(i + 1) for i in range(args.count)]))

    client.terminate()

    print("DONE", file=sys.stderr)  # noqa: T201
