#
# 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/.

""" Define Qgis Worker

    We define QGis workers as DEALER configuration: it will enables
    us to forward streamed response as chunked transfer.

    1. Start qgis application
    2. Connect a REQ zmq socket to endpoint
    3. Return qgis server responses
"""
import logging
import os
import pickle
import sys
import traceback
import uuid

from typing import (
    Callable,
    Dict,
    Mapping,
    Optional,
    Protocol,
    Type,
    Union,
)

import zmq

from ..logger import setup_log_handler
from ..utils import stats
from .messages import WORKER_READY, ReplyMessage
from .supervisor import Client as SupervisorClient

LOGGER = logging.getLogger('SRVLOG')


# Define an abstract type for HTTPRequest
class HTTPRequest(Protocol):
    headers: Mapping[str, str]
    query: str
    data: bytes
    method: str


class RequestHandler:

    def __init__(
        self,
        socket: zmq.Socket,
        client_id: bytes,
        correlation_id: bytes,
        request: HTTPRequest,
    ):
        """ Handle requests and

            Handle reply message contruction and pass message correlation_id to
            reply.

            :param request: An HTTP request handler
            :param socket: the zmq socket
        """
        self.headers: Dict = {}
        self.status_code = 200
        self.request = request
        self.header_written = False

        self._correlation_id = correlation_id
        self._socket = socket
        self._client_id = client_id

    def _write(self, data: bytes):
        """ Send data back to client
        """
        self._socket.send_multipart([
            self._client_id,
            self._correlation_id,
            data])

    def send(self, data: bytes, send_more: bool = False, meta: Optional[Dict[str, str]] = None):
        """ Send data
        """
        if not self.header_written:
            # We let the client know that there is more
            # data by setting the 206 code (partial response)
            if send_more and self.status_code == 200:
                self.status_code = 206
            # Create a Header Message
            message = pickle.dumps(
                ReplyMessage(
                    self.status_code,
                    headers=self.headers,
                    data=data,
                    meta=meta,
                ),
                -1,
            )
            self._write(message)
            self.header_written = True
        elif self.status_code == 206:
            self._write(pickle.dumps((data, send_more, meta), -1))
            if not send_more:
                self.status_code = 200

    @property
    def msgid(self):
        """ Return the message correlation id
        """
        return self._correlation_id

    @property
    def identity(self) -> Union[bytes, str, int]:
        return self._socket.identity

    def handle_message(self):
        """ Override this method to handle_messages
        """
        print("Received", self.request.data, "from", self._client_id)  # noqa T201
        self.send(b"Hello %s" % self._client_id, True)
        self.send(b"Chunk 1", True)
        self.send(b"Chunk 2", True)
        self.send(b"", False)

    @classmethod
    def get_report(cls):
        data = stats.stats()
        data.update(pid=os.getpid())
        return data


def dealer_socket(ctx: zmq.Context, address: str, identity: Optional[bytes] = None) -> zmq.Socket:
    """ Socket for receiving incoming messages
    """
    LOGGER.debug("Connecting to %s", address)
    sock = ctx.socket(zmq.DEALER)
    sock.setsockopt(zmq.LINGER, 500)     # Needed for socket no to wait on close
    sock.setsockopt(zmq.SNDHWM, 1)       # Max 1 item on send queue
    sock.setsockopt(zmq.IMMEDIATE, 1)    # Do not queue if no peer, will block on send
    sock.setsockopt(zmq.RCVTIMEO, 1000)  # Heartbeat
    sock.identity = identity or uuid.uuid1().bytes
    LOGGER.debug("Identity set to %s", sock.identity)
    sock.connect(address)
    return sock


def broadcast_socket(ctx: zmq.Context, broadcastaddr: str) -> zmq.Socket:
    """ Socket for receiving broadcast message notifications
    """
    LOGGER.debug("Enabling broadcast notification")
    ctx = zmq.Context.instance()
    sub = ctx.socket(zmq.SUB)
    sub.setsockopt(zmq.LINGER, 500)    # Needed for socket no to wait on close
    sub.setsockopt(zmq.SUBSCRIBE, b'RESTART')
    sub.setsockopt(zmq.SUBSCRIBE, b'REPORT')
    sub.connect(broadcastaddr)
    return sub


def run_worker(
    address: str,
    handler_factory: Type[RequestHandler],
    identity: Optional[bytes] = None,
    broadcastaddr: Optional[str] = None,
    postprocess: Optional[Callable[[bool], None]] = None,
):
    """ Enter the message loop
    """
    ctx = zmq.Context.instance()

    sock = dealer_socket(ctx, address, identity)
    if broadcastaddr:
        sub = broadcast_socket(ctx, broadcastaddr)

    # Initialize supervisor client
    supervisor = SupervisorClient()

    def get():
        client_id, corr_id, request = sock.recv_multipart()
        LOGGER.debug("RCV %s: %s", client_id, corr_id)
        return client_id, corr_id, pickle.loads(request)

    try:
        LOGGER.info("Starting ZMQ worker loop")
        while True:
            sock.send(WORKER_READY)
            idle = False
            handler = None
            try:
                client_id, corr_id, request = get()
                supervisor.notify_busy()
                handler = handler_factory(sock, client_id, corr_id, request)
                handler.handle_message()
            except zmq.error.Again:
                idle = True
                pass
            except zmq.ZMQError as err:
                LOGGER.error("Worker Error %d: %s", err.errno, zmq.strerror(err.errno))
            except Exception as exc:
                LOGGER.error("Worker Error %s", exc)
                # Print trace outside LOGGER because
                # logging output in sub-processe
                # is not captured by pytest
                traceback.print_exc()
                if handler and not handler.header_written:
                    handler.status_code = 500
                    handler.send(b"Worker internal error")
                    # Got error 500, do not presume worker state
                    break
            finally:
                supervisor.notify_done()

            # Handle broadcast notifications
            try:
                if broadcastaddr:
                    msg = sub.recv(flags=zmq.NOBLOCK)
                    if msg == b'RESTART':
                        # There is no really way to restart
                        # so exit and let the framework restart a new worker
                        LOGGER.info("Exiting on RESTART notification")
                        break
                    elif msg == b'REPORT':
                        # Reporting asked
                        supervisor.send_report(handler_factory.get_report())
            except zmq.error.Again:
                pass

            try:
                # Run callbacks
                if postprocess:
                    postprocess(idle)
            except Exception:
                LOGGER.critical("Unhandled exception:\n%s", traceback.format_exc())

    except (KeyboardInterrupt, SystemExit):
        pass

    if broadcastaddr:
        sub.close()
    sock.close()
    LOGGER.info("Terminating Worker")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Test worker')
    parser.add_argument('--host', metavar="host", default='tcp://localhost', help="router host")
    parser.add_argument('--router', metavar='address', default='{host}:8881', 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")

    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()))

    run_worker(
        args.router.format(host=args.host),
        RequestHandler,
        identity=bytes(args.identity.encode('ascii')),
    )

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