import asyncio
import logging
import multiprocessing
import os
import socket
import time
import warnings

import aiohttp
import aiohttp.web
import tornado.ioloop

import salt.payload
import salt.transport.base
import salt.transport.frame
import salt.utils.asynchronous
from salt.transport.tcp import (
    USE_LOAD_BALANCER,
    LoadBalancerServer,
    _get_bind_addr,
    _get_socket,
    _set_tcp_keepalive,
)

log = logging.getLogger(__name__)


class PublishClient(salt.transport.base.PublishClient):
    """
    Tornado based TCP Pub Client
    """

    ttype = "ws"

    async_methods = [
        "connect",
        "connect_uri",
        "recv",
        # "close",
    ]
    close_methods = [
        "close",
    ]

    def __init__(self, opts, io_loop, **kwargs):  # pylint: disable=W0231
        self.opts = opts
        if io_loop is None:
            io_loop = tornado.ioloop.IOLoop.current()
        self.io_loop = io_loop
        self.asyncio_loop = salt.utils.asynchronous.aioloop(io_loop)

        self.connected = False
        self._closing = False
        self._closing = False
        self._closed = False

        self.backoff = opts.get("tcp_reconnect_backoff", 1)
        self.poller = None

        self.host = kwargs.get("host", None)
        self.port = kwargs.get("port", None)
        self.path = kwargs.get("path", None)
        self.ssl = kwargs.get("ssl", None)
        self.source_ip = self.opts.get("source_ip")
        self.source_port = self.opts.get("source_publish_port")
        self.connect_callback = None
        self.disconnect_callback = None
        if self.host is None and self.port is None:
            if self.path is None:
                raise Exception("A host and port or a path must be provided")
        elif self.host and self.port:
            if self.path:
                raise Exception("A host and port or a path must be provided, not both")
        self._ws = None
        self._session = None
        self._closing = False
        self.on_recv_task = None

    def close(self):
        if self._closing:
            return
        self._closing = True
        # Cancel the receive task but don't await it (like TCP does)
        if self.on_recv_task:
            self.on_recv_task.cancel()
            self.on_recv_task = None
        # Schedule async cleanup but don't wait for it
        if self._session is not None or self._ws is not None:
            self.asyncio_loop.create_task(self._async_cleanup())
        self._closed = True

    async def _async_cleanup(self):
        """Background cleanup of async resources"""
        try:
            if self._session is not None:
                await self._session.close()
                self._session = None
            if self._ws is not None:
                await self._ws.close()
                self._ws = None
        except Exception:  # pylint: disable=broad-except
            pass  # Cleanup is best-effort

    # pylint: disable=W1701
    def __del__(self):
        if not self._closing:
            warnings.warn(
                "unclosed publish client {self!r}", ResourceWarning, source=self
            )

    # pylint: enable=W1701
    def _decode_messages(self, messages):
        if not isinstance(messages, dict):
            body = salt.payload.loads(messages)
        else:
            body = messages
        return body

    async def getstream(self, **kwargs):
        if self.source_ip or self.source_port:
            kwargs.update(source_ip=self.source_ip, source_port=self.source_port)
        ws = None
        session = None
        start = time.monotonic()
        timeout = kwargs.get("timeout", None)
        while ws is None and (not self._closed and not self._closing):
            try:
                ctx = None
                if self.ssl is not None:
                    ctx = salt.transport.base.ssl_context(self.ssl, server_side=False)
                if self.host and self.port:
                    conn = aiohttp.TCPConnector()
                    session = aiohttp.ClientSession(connector=conn)
                    if self.ssl:
                        url = f"https://{self.host}:{self.port}/ws"
                    else:
                        url = f"http://{self.host}:{self.port}/ws"
                else:
                    conn = aiohttp.UnixConnector(path=self.path)
                    session = aiohttp.ClientSession(connector=conn)
                    if self.ssl:
                        url = "https://ipc.saltproject.io/ws"
                    else:
                        url = "http://ipc.saltproject.io/ws"
                log.debug("pub client connect %r %r", url, ctx)
                ws = await asyncio.wait_for(
                    session.ws_connect(url, ssl=ctx),
                    timeout if timeout is not None else 5,
                )
                # For SSL connections, give handshake time to complete and fail if invalid
                if ws and self.ssl:
                    await asyncio.sleep(0.1)
                    if ws.closed:
                        log.debug("WS closed after SSL handshake")
                        ws = None
                        await session.close()
                        continue
            except Exception as exc:  # pylint: disable=broad-except
                log.warning(
                    "WS Message Client encountered an exception while connecting to"
                    " %s:%s %s: %r, will reconnect in %d seconds",
                    self.host,
                    self.port,
                    self.path,
                    exc,
                    self.backoff,
                )
                if session:
                    await session.close()
                if timeout and time.monotonic() - start > timeout:
                    break
                await asyncio.sleep(self.backoff)
        return ws, session

    async def _connect(self, timeout=None):
        if self._ws is None:
            self._ws, self._session = await self.getstream(timeout=timeout)
            if self.connect_callback:
                self.connect_callback(True)  # pylint: disable=not-callable
            self.connected = True

    async def connect(
        self,
        port=None,
        connect_callback=None,
        disconnect_callback=None,
        timeout=None,
    ):
        if port is not None:
            self.port = port
        if connect_callback:
            self.connect_callback = None
        if disconnect_callback:
            self.disconnect_callback = None
        await self._connect(timeout=timeout)

    async def send(self, msg):
        await self.message_client.send(msg, reply=False)

    async def recv(self, timeout=None):
        while self._ws is None:
            await self.connect()
            await asyncio.sleep(0.001)
        if timeout == 0:
            try:
                raw_msg = await asyncio.wait_for(self._ws.receive(), 0.0001)
            except TimeoutError:
                return
            if raw_msg.type == aiohttp.WSMsgType.TEXT:
                if raw_msg.data == "close":
                    await self._ws.close()
            if raw_msg.type == aiohttp.WSMsgType.BINARY:
                return salt.payload.loads(raw_msg.data, raw=True)
            elif raw_msg.type == aiohttp.WSMsgType.ERROR:
                log.error(
                    "ws connection closed with exception %s", self._ws.exception()
                )
        elif timeout:
            return await asyncio.wait_for(self.recv(), timeout=timeout)
        else:
            while True:
                raw_msg = await self._ws.receive()
                if raw_msg.type == aiohttp.WSMsgType.TEXT:
                    if raw_msg.data == "close":
                        await self._ws.close()
                if raw_msg.type == aiohttp.WSMsgType.BINARY:
                    return salt.payload.loads(raw_msg.data, raw=True)
                elif raw_msg.type == aiohttp.WSMsgType.ERROR:
                    log.error(
                        "ws connection closed with exception %s",
                        self._ws.exception(),
                    )

    async def on_recv_handler(self, callback):
        while not self._ws:
            await asyncio.sleep(0.003)
        while True:
            msg = await self.recv()
            await callback(msg)

    def on_recv(self, callback):
        """
        Register a callback for received messages (that we didn't initiate)
        """
        if self.on_recv_task:
            # XXX: We are not awaiting this canceled task. This still needs to
            # be addressed.
            self.on_recv_task.cancel()
        if callback is None:
            self.on_recv_task = None
        else:
            self.on_recv_task = asyncio.create_task(self.on_recv_handler(callback))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class PublishServer(salt.transport.base.DaemonizedPublishServer):
    """ """

    # TODO: opts!
    # Based on default used in tornado.netutil.bind_sockets()
    backlog = 128
    async_methods = [
        "publish",
        # "close",
    ]
    close_methods = [
        "close",
    ]

    def __init__(
        self,
        opts,
        pub_host=None,
        pub_port=None,
        pub_path=None,
        pull_host=None,
        pull_port=None,
        pull_path=None,
        pull_path_perms=0o600,
        pub_path_perms=0o600,
        started=None,
        _shutdown=None,
        ssl=None,
    ):
        self.opts = opts
        self.pub_host = pub_host
        self.pub_port = pub_port
        self.pub_path = pub_path
        self.pull_host = pull_host
        self.pull_port = pull_port
        self.pull_path = pull_path
        self.pull_path_perms = pull_path_perms
        self.pub_path_perms = pub_path_perms
        self.ssl = ssl
        self.clients = set()
        self._run = None
        if _shutdown is None:
            self._shutdown = multiprocessing.Event()  # Cross-process shutdown signal
        else:
            self._shutdown = _shutdown
        self.pub_writer = None
        self.pub_reader = None
        self._connecting = None
        self.runner = None
        self.site = None
        self.puller = None
        if started is None:
            self.started = multiprocessing.Event()
        else:
            self.started = started

    @classmethod
    def support_ssl(cls):
        # Required from DaemonizedPublishServer
        return True

    def topic_support(self):
        # Required from DaemonizedPublishServer
        return not self.opts.get("order_masters", False)

    def __setstate__(self, state):
        self.__init__(**state)

    def __getstate__(self):
        return {
            "opts": self.opts,
            "pub_host": self.pub_host,
            "pub_port": self.pub_port,
            "pub_path": self.pub_path,
            "pull_host": self.pull_host,
            "pull_port": self.pull_port,
            "pull_path": self.pull_path,
            "pull_path_perms": self.pull_path_perms,
            "pub_path_perms": self.pub_path_perms,
            "ssl": self.ssl,
            "started": self.started,
            "_shutdown": self._shutdown,
        }

    def publish_daemon(
        self,
        publish_payload,
        presence_callback=None,
        remove_presence_callback=None,
        secrets=None,
        started=None,
    ):
        """
        Bind to the interface specified in the configuration file
        """
        if started is not None:
            self.started = started
        # Use asyncio event loop directly like ZeroMQ does
        io_loop = salt.utils.asynchronous.aioloop(tornado.ioloop.IOLoop())

        # Set up asyncio signal handler to stop the loop on SIGTERM
        import signal

        io_loop.add_signal_handler(signal.SIGTERM, io_loop.stop)

        publisher_task = io_loop.create_task(
            self.publisher(publish_payload, io_loop=io_loop)
        )
        try:
            io_loop.run_forever()
        finally:
            self.close()

    async def publisher(
        self,
        publish_payload,
        presence_callback=None,
        remove_presence_callback=None,
        io_loop=None,
    ):
        if io_loop is None:
            io_loop = tornado.ioloop.IOLoop.current()
        if self._run is None:
            self._run = asyncio.Event()

        # Monitor the multiprocessing shutdown event and stop the loop
        async def monitor_shutdown():
            loop = asyncio.get_running_loop()
            # Wait for shutdown signal in executor to avoid blocking
            await loop.run_in_executor(None, self._shutdown.wait)
            self._run.set()
            loop.stop()

        asyncio.create_task(monitor_shutdown())

        ctx = None
        if self.ssl is not None:
            ctx = salt.transport.base.ssl_context(self.ssl, server_side=True)
        if self.pub_path:
            server = aiohttp.web.Server(self.handle_request)
            self.runner = aiohttp.web.ServerRunner(server)
            await self.runner.setup()
            with salt.utils.files.set_umask(0o177):
                log.info("Publisher binding to socket %s", self.pub_path)
                self.site = aiohttp.web.UnixSite(
                    self.runner, self.pub_path, ssl_context=ctx
                )
                await self.site.start()
                os.chmod(self.pub_path, self.pub_path_perms)
        else:
            sock = _get_socket(self.opts)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            # _set_tcp_keepalive(sock, self.opts)
            sock.setblocking(0)
            sock.bind((self.pub_host, self.pub_port))
            sock.listen(self.backlog)
            server = aiohttp.web.Server(self.handle_request)
            self.runner = aiohttp.web.ServerRunner(server)
            await self.runner.setup()
            self.site = aiohttp.web.SockSite(self.runner, sock, ssl_context=ctx)
            log.info("Publisher binding to socket %s:%s", self.pub_host, self.pub_port)
            await self.site.start()

        self._pub_payload = publish_payload
        if self.pull_path:
            with salt.utils.files.set_umask(0o177):
                self.puller = await asyncio.start_unix_server(
                    self.pull_handler, self.pull_path
                )
                os.chmod(self.pull_path, self.pull_path_perms)
        else:
            self.puller = await asyncio.start_server(
                self.pull_handler, self.pull_host, self.pull_port
            )
        self.started.set()
        # Wait for shutdown signal instead of polling
        await self._run.wait()
        # Properly shut down aiohttp server
        await self.site.stop()
        await self.runner.cleanup()
        # Close the puller server
        self.puller.close()
        await self.puller.wait_closed()

    async def pull_handler(self, reader, writer):
        unpacker = salt.utils.msgpack.Unpacker(raw=True)
        while True:
            data = await reader.read(1024)
            unpacker.feed(data)
            for msg in unpacker:
                await self._pub_payload(msg)

    def pre_fork(self, process_manager, *args, **kwargs):
        """
        Do anything necessary pre-fork. Since this is on the master side this will
        primarily be used to create IPC channels and create our daemon process to
        do the actual publishing
        """
        process_manager.add_process(
            self.publish_daemon,
            args=[self.publish_payload],
            name=self.__class__.__name__,
        )

    async def handle_request(self, request):
        try:
            cert = request.get_extra_info("peercert")
        except AttributeError:
            pass
        else:
            if cert:
                name = salt.transport.base.common_name(cert)
                log.debug("Request client cert %r", name)
        ws = aiohttp.web.WebSocketResponse()
        await ws.prepare(request)
        self.clients.add(ws)
        try:
            # Keep connection alive until client disconnects
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.ERROR:
                    log.error("ws connection closed with exception %s", ws.exception())
                    break
        finally:
            self.clients.discard(ws)
        return ws

    async def _connect(self):
        if self.pull_path:
            self.pub_reader, self.pub_writer = await asyncio.open_unix_connection(
                self.pull_path
            )
        else:
            self.pub_reader, self.pub_writer = await asyncio.open_connection(
                self.pull_host, self.pull_port
            )
        self._connecting = None

    def connect(self, timeout=None):
        log.debug("Connect pusher %s", self.pull_path)
        if self._connecting is None:
            self._connecting = asyncio.create_task(self._connect())
        return self._connecting

    async def publish(
        self, payload, **kwargs
    ):  # pylint: disable=invalid-overridden-method
        """
        Publish "load" to minions
        """
        if not self.pub_writer:
            await self.connect()
        self.pub_writer.write(salt.payload.dumps(payload, use_bin_type=True))
        await self.pub_writer.drain()

    async def publish_payload(self, payload, topic_list=None, raw_payload=None):
        # ``raw_payload`` is accepted for interface parity with the
        # TCP PublishServer, which uses it to skip a redundant
        # msgpack round-trip on the EP fan-out hot path.  The ws
        # transport frames with ``salt.payload.dumps`` below, so the
        # unframed passthrough shortcut doesn't apply here.
        payload = salt.payload.dumps(payload, use_bin_type=True)
        for ws in list(self.clients):
            try:
                await ws.send_bytes(payload)
            except ConnectionResetError:
                self.clients.discard(ws)

    def close(self):
        if self.pub_writer:
            self.pub_writer.close()
            self.pub_writer = None
            self.pub_reader = None
        # Signal shutdown across processes
        self._shutdown.set()
        if self._connecting:
            self._connecting.cancel()


class RequestServer(salt.transport.base.DaemonizedRequestServer):
    def __init__(self, opts):  # pylint: disable=W0231
        self.opts = opts
        self.site = None
        self.ssl = self.opts.get("ssl", None)
        self._run = None
        self._socket = None

    def pre_fork(self, process_manager, *args, **kwargs):
        """
        Pre-fork we need to create the zmq router device
        """
        if USE_LOAD_BALANCER:
            self.socket_queue = multiprocessing.Queue()
            process_manager.add_process(
                LoadBalancerServer,
                args=(self.opts, self.socket_queue),
                name="LoadBalancerServer",
            )
        elif not salt.utils.platform.is_windows():
            if self.opts.get("ipc_mode") == "ipc" and self.opts.get("workers_ipc_name"):
                self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                self._socket.setblocking(0)
                ipc_path = os.path.join(
                    self.opts["sock_dir"], self.opts["workers_ipc_name"]
                )
                if os.path.exists(ipc_path):
                    os.unlink(ipc_path)
                self._socket.bind(ipc_path)
                os.chmod(ipc_path, 0o600)
            else:
                self._socket = _get_socket(self.opts)
                self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                _set_tcp_keepalive(self._socket, self.opts)
                self._socket.setblocking(0)
                self._socket.bind(_get_bind_addr(self.opts, "ret_port"))

    def post_fork(self, message_handler, io_loop, **kwargs):
        """
        After forking we need to create all of the local sockets to listen to the
        router

        message_handler: function to call with your payloads
        """
        self.message_handler = message_handler
        self._run = asyncio.Event()
        self._started = asyncio.Event()

        # Convert to asyncio loop
        io_loop = salt.utils.asynchronous.aioloop(io_loop)

        async def server():
            server = aiohttp.web.Server(self.handle_message)
            runner = aiohttp.web.ServerRunner(server)
            await runner.setup()
            ctx = None
            if self.ssl is not None:
                ctx = salt.transport.base.ssl_context(self.ssl, server_side=True)
            self.site = aiohttp.web.SockSite(runner, self._socket, ssl_context=ctx)
            log.info("Worker binding to socket %s", self._socket)
            await self.site.start()
            self._started.set()
            # pause here for very long time by serving HTTP requests and
            # waiting for keyboard interruption
            # Wait for shutdown signal instead of polling
            await self._run.wait()
            await self.site.stop()
            self._socket.close()

        io_loop.create_task(server())

    async def handle_message(self, request):
        try:
            cert = request.get_extra_info("peercert")
        except AttributeError:
            pass
        else:
            if cert:
                name = salt.transport.base.common_name(cert)
                log.debug("Request client cert %r", name)
        ws = aiohttp.web.WebSocketResponse()
        await ws.prepare(request)
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                if msg.data == "close":
                    await ws.close()
            if msg.type == aiohttp.WSMsgType.BINARY:
                payload = salt.payload.loads(msg.data)
                reply = await self.message_handler(payload)
                await ws.send_bytes(salt.payload.dumps(reply))
            elif msg.type == aiohttp.WSMsgType.ERROR:
                log.error("ws connection closed with exception %s", ws.exception())
        return ws

    def close(self):
        if self._run is not None:
            self._run.set()  # Signal shutdown
        if self._socket is not None:
            self._socket.shutdown(socket.SHUT_RDWR)
            self._socket.close()
            self._socket = None

    async def forward_message(self, payload):
        """
        Forward a message into this transport's worker queue.

        Not implemented for WebSocket transport. Worker pool routing is only
        supported for ZeroMQ transport.
        """
        log.warning(
            "Worker pool message forwarding is not supported for WebSocket transport. "
            "Use ZeroMQ transport for worker pool routing."
        )
        return None


class RequestClient(salt.transport.base.RequestClient):

    ttype = "ws"

    def __init__(self, opts, io_loop):  # pylint: disable=W0231
        self.opts = opts
        self.sending = False
        self.ws = None
        self.session = None
        self.io_loop = salt.utils.asynchronous.aioloop(io_loop)
        self._closing = False
        self._closed = False
        self.ssl = self.opts.get("ssl", None)

    async def connect(self):  # pylint: disable=invalid-overridden-method
        ctx = None
        if self.ssl is not None:
            ctx = salt.transport.base.ssl_context(self.ssl, server_side=False)

        master_uri = self.opts.get("master_uri", "")
        if master_uri.startswith("ipc://"):
            socket_path = master_uri[6:]
            connector = aiohttp.UnixConnector(path=socket_path)
            self.session = aiohttp.ClientSession(connector=connector)
            URL = "http://localhost/ws"
        else:
            self.session = aiohttp.ClientSession()
            URL = self.get_master_uri(self.opts)

        log.debug("Connect to %s %s", URL, ctx)
        self.ws = await self.session.ws_connect(URL, ssl=ctx)

    async def send(self, load, timeout=60):
        if self.sending or self._closing:
            await asyncio.sleep(0.03)
        self.sending = True
        try:
            await self.connect()
            message = salt.payload.dumps(load)
            await self.ws.send_bytes(message)
            async for msg in self.ws:
                if msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break
                data = salt.payload.loads(msg.data)
                break
            return data
        finally:
            self.sending = False

    async def _close(self):
        if self.ws is not None:
            await self.ws.close()
            self.ws = None
        if self.session is not None:
            await self.session.close()
            self.session = None
        self._closed = True

    def close(self):
        if self._closing:
            return
        self._closing = True
        self.close_task = asyncio.create_task(self._close())

    def get_master_uri(self, opts):
        if "master_uri" in opts:
            if self.opts.get("ssl", None):
                return opts["master_uri"].replace("tcp:", "https:", 1)
            return opts["master_uri"].replace("tcp:", "http:", 1)
        if self.opts.get("ssl", None):
            return f"https://{opts['master_ip']}:{opts['master_port']}/ws"
        return f"http://{opts['master_ip']}:{opts['master_port']}/ws"

    # pylint: disable=W1701
    def __del__(self):
        if not self._closing:
            warnings.warn(
                "Unclosed publish client {self!r}", ResourceWarning, source=self
            )

    # pylint: enable=W1701
