from __future__ import print_function, division, absolute_import

import errno
import logging
import socket
import ssl
import struct
import sys

import tornado
from tornado import gen, netutil
from tornado.iostream import StreamClosedError
from tornado.tcpclient import TCPClient
from tornado.tcpserver import TCPServer

from .. import config
from ..compatibility import finalize
from ..utils import (ensure_bytes, ensure_ip, get_ip, get_ipv6, nbytes)

from .registry import Backend, backends
from .addressing import parse_host_port, unparse_host_port
from .core import Comm, Connector, Listener, CommClosedError
from .utils import (to_frames, from_frames,
                    get_tcp_server_address, ensure_concrete_host)


logger = logging.getLogger(__name__)


def get_total_physical_memory():
    try:
        import psutil
        return psutil.virtual_memory().total / 2
    except ImportError:
        return 2e9


MAX_BUFFER_SIZE = get_total_physical_memory()

DEFAULT_BACKLOG = 2048


def set_tcp_timeout(stream):
    """
    Set kernel-level TCP timeout on the stream.
    """
    if stream.closed():
        return

    timeout = int(config.get('tcp-timeout', 30))

    sock = stream.socket

    # Default (unsettable) value on Windows
    # https://msdn.microsoft.com/en-us/library/windows/desktop/dd877220(v=vs.85).aspx
    nprobes = 10
    assert timeout >= nprobes + 1, "Timeout too low"

    idle = max(2, timeout // 4)
    interval = max(1, (timeout - idle) // nprobes)
    idle = timeout - interval * nprobes
    assert idle > 0

    try:
        if sys.platform.startswith("win"):
            logger.debug("Setting TCP keepalive: idle=%d, interval=%d",
                         idle, interval)
            sock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, idle * 1000, interval * 1000))
        else:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            try:
                TCP_KEEPIDLE = socket.TCP_KEEPIDLE
                TCP_KEEPINTVL = socket.TCP_KEEPINTVL
                TCP_KEEPCNT = socket.TCP_KEEPCNT
            except AttributeError:
                if sys.platform == "darwin":
                    TCP_KEEPIDLE = 0x10  # (named "TCP_KEEPALIVE" in C)
                    TCP_KEEPINTVL = 0x101
                    TCP_KEEPCNT = 0x102
                else:
                    TCP_KEEPIDLE = None

            if TCP_KEEPIDLE is not None:
                logger.debug("Setting TCP keepalive: nprobes=%d, idle=%d, interval=%d",
                             nprobes, idle, interval)
                sock.setsockopt(socket.SOL_TCP, TCP_KEEPCNT, nprobes)
                sock.setsockopt(socket.SOL_TCP, TCP_KEEPIDLE, idle)
                sock.setsockopt(socket.SOL_TCP, TCP_KEEPINTVL, interval)

        if sys.platform.startswith("linux"):
            logger.debug("Setting TCP user timeout: %d ms",
                         timeout * 1000)
            TCP_USER_TIMEOUT = 18  # since Linux 2.6.37
            sock.setsockopt(socket.SOL_TCP, TCP_USER_TIMEOUT, timeout * 1000)
    except EnvironmentError as e:
        logger.warning("Could not set timeout on TCP stream: %s", e)


def get_stream_address(stream):
    """
    Get a stream's local address.
    """
    if stream.closed():
        return "<closed>"

    try:
        return unparse_host_port(*stream.socket.getsockname()[:2])
    except EnvironmentError:
        # Probably EBADF
        return "<closed>"


def convert_stream_closed_error(obj, exc):
    """
    Re-raise StreamClosedError as CommClosedError.
    """
    if exc.real_error is not None:
        # The stream was closed because of an underlying OS error
        exc = exc.real_error
        raise CommClosedError("in %s: %s: %s" % (obj, exc.__class__.__name__, exc))
    else:
        raise CommClosedError("in %s: %s" % (obj, exc))


class TCP(Comm):
    """
    An established communication based on an underlying Tornado IOStream.
    """
    _iostream_allows_memoryview = tornado.version_info >= (4, 5)

    def __init__(self, stream, local_addr, peer_addr, deserialize=True):
        self._local_addr = local_addr
        self._peer_addr = peer_addr
        self.stream = stream
        self.deserialize = deserialize
        self._finalizer = finalize(self, self._get_finalizer())
        self._finalizer.atexit = False
        self._extra = {}

        stream.set_nodelay(True)
        set_tcp_timeout(stream)
        self._read_extra()

    def _read_extra(self):
        pass

    def _get_finalizer(self):
        def finalize(stream=self.stream, r=repr(self)):
            if not stream.closed():
                logger.warning("Closing dangling stream in %s" % (r,))
                stream.close()

        return finalize

    @property
    def local_address(self):
        return self._local_addr

    @property
    def peer_address(self):
        return self._peer_addr

    @gen.coroutine
    def read(self):
        stream = self.stream
        if stream is None:
            raise CommClosedError

        try:
            n_frames = yield stream.read_bytes(8)
            n_frames = struct.unpack('Q', n_frames)[0]
            lengths = yield stream.read_bytes(8 * n_frames)
            lengths = struct.unpack('Q' * n_frames, lengths)

            frames = []
            for length in lengths:
                if length:
                    frame = yield stream.read_bytes(length)
                else:
                    frame = b''
                frames.append(frame)
        except StreamClosedError as e:
            self.stream = None
            convert_stream_closed_error(self, e)

        try:
            msg = yield from_frames(frames, deserialize=self.deserialize)
        except EOFError:
            # Frames possibly garbled or truncated by communication error
            self.abort()
            raise CommClosedError("aborted stream on truncated data")
        raise gen.Return(msg)

    @gen.coroutine
    def write(self, msg):
        stream = self.stream
        bytes_since_last_yield = 0
        if stream is None:
            raise CommClosedError

        frames = yield to_frames(msg)

        try:
            lengths = ([struct.pack('Q', len(frames))] +
                       [struct.pack('Q', nbytes(frame)) for frame in frames])
            stream.write(b''.join(lengths))

            for frame in frames:
                # Can't wait for the write() Future as it may be lost
                # ("If write is called again before that Future has resolved,
                #   the previous future will be orphaned and will never resolve")
                if not self._iostream_allows_memoryview:
                    frame = ensure_bytes(frame)
                future = stream.write(frame)
                bytes_since_last_yield += nbytes(frame)
                if bytes_since_last_yield > 32e6:
                    yield future
                    bytes_since_last_yield = 0
        except StreamClosedError as e:
            stream = None
            convert_stream_closed_error(self, e)
        except TypeError as e:
            if stream._write_buffer is None:
                logger.info("tried to write message %s on closed stream", msg)
            else:
                raise

        raise gen.Return(sum(map(nbytes, frames)))

    @gen.coroutine
    def close(self):
        stream, self.stream = self.stream, None
        if stream is not None and not stream.closed():
            try:
                # Flush the stream's write buffer by waiting for a last write.
                if stream.writing():
                    yield stream.write(b'')
                stream.socket.shutdown(socket.SHUT_RDWR)
            except EnvironmentError:
                pass
            finally:
                self._finalizer.detach()
                stream.close()

    def abort(self):
        stream, self.stream = self.stream, None
        if stream is not None and not stream.closed():
            self._finalizer.detach()
            stream.close()

    def closed(self):
        return self.stream is None or self.stream.closed()

    @property
    def extra_info(self):
        return self._extra


class TLS(TCP):
    """
    A TLS-specific version of TCP.
    """

    def _read_extra(self):
        TCP._read_extra(self)
        sock = self.stream.socket
        if sock is not None:
            self._extra.update(peercert=sock.getpeercert(),
                               cipher=sock.cipher())
            cipher, proto, bits = self._extra['cipher']
            logger.debug("TLS connection with %r: protocol=%s, cipher=%s, bits=%d",
                         self._peer_addr, proto, cipher, bits)


def _expect_tls_context(connection_args):
    ctx = connection_args.get('ssl_context')
    if not isinstance(ctx, ssl.SSLContext):
        raise TypeError("TLS expects a `ssl_context` argument of type "
                        "ssl.SSLContext (perhaps check your TLS configuration?)")
    return ctx


class RequireEncryptionMixin(object):

    def _check_encryption(self, address, connection_args):
        if not self.encrypted and connection_args.get('require_encryption'):
            # XXX Should we have a dedicated SecurityError class?
            raise RuntimeError("encryption required by Dask configuration, "
                               "refusing communication from/to %r"
                               % (self.prefix + address,))


class BaseTCPConnector(Connector, RequireEncryptionMixin):

    @gen.coroutine
    def connect(self, address, deserialize=True, **connection_args):
        self._check_encryption(address, connection_args)
        ip, port = parse_host_port(address)
        kwargs = self._get_connect_args(**connection_args)

        client = TCPClient()
        try:
            stream = yield client.connect(ip, port,
                                          max_buffer_size=MAX_BUFFER_SIZE,
                                          **kwargs)
        except StreamClosedError as e:
            # The socket connect() call failed
            convert_stream_closed_error(self, e)

        local_address = self.prefix + get_stream_address(stream)
        raise gen.Return(self.comm_class(stream,
                                         local_address,
                                         self.prefix + address,
                                         deserialize))


class TCPConnector(BaseTCPConnector):
    prefix = 'tcp://'
    comm_class = TCP
    encrypted = False

    def _get_connect_args(self, **connection_args):
        return {}


class TLSConnector(BaseTCPConnector):
    prefix = 'tls://'
    comm_class = TLS
    encrypted = True

    def _get_connect_args(self, **connection_args):
        ctx = _expect_tls_context(connection_args)
        return {'ssl_options': ctx}


class BaseTCPListener(Listener, RequireEncryptionMixin):

    def __init__(self, address, comm_handler, deserialize=True,
                 default_port=0, **connection_args):
        self._check_encryption(address, connection_args)
        self.ip, self.port = parse_host_port(address, default_port)
        self.comm_handler = comm_handler
        self.deserialize = deserialize
        self.server_args = self._get_server_args(**connection_args)
        self.tcp_server = None
        self.bound_address = None

    def start(self):
        self.tcp_server = TCPServer(max_buffer_size=MAX_BUFFER_SIZE,
                                    **self.server_args)
        self.tcp_server.handle_stream = self._handle_stream
        backlog = int(config.get('socket-backlog', DEFAULT_BACKLOG))
        for i in range(5):
            try:
                # When shuffling data between workers, there can
                # really be O(cluster size) connection requests
                # on a single worker socket, make sure the backlog
                # is large enough not to lose any.
                sockets = netutil.bind_sockets(self.port, address=self.ip,
                                               backlog=backlog)
            except EnvironmentError as e:
                # EADDRINUSE can happen sporadically when trying to bind
                # to an ephemeral port
                if self.port != 0 or e.errno != errno.EADDRINUSE:
                    raise
                exc = e
            else:
                self.tcp_server.add_sockets(sockets)
                break
        else:
            raise exc

    def stop(self):
        tcp_server, self.tcp_server = self.tcp_server, None
        if tcp_server is not None:
            tcp_server.stop()

    def _check_started(self):
        if self.tcp_server is None:
            raise ValueError("invalid operation on non-started TCPListener")

    @gen.coroutine
    def _handle_stream(self, stream, address):
        address = self.prefix + unparse_host_port(*address[:2])
        stream = yield self._prepare_stream(stream, address)
        if stream is None:
            # Preparation failed
            return
        logger.debug("Incoming connection from %r to %r",
                     address, self.contact_address)
        local_address = self.prefix + get_stream_address(stream)
        comm = self.comm_class(stream, local_address, address, self.deserialize)
        self.comm_handler(comm)

    def get_host_port(self):
        """
        The listening address as a (host, port) tuple.
        """
        self._check_started()

        if self.bound_address is None:
            self.bound_address = get_tcp_server_address(self.tcp_server)
        # IPv6 getsockname() can return more a 4-len tuple
        return self.bound_address[:2]

    @property
    def listen_address(self):
        """
        The listening address as a string.
        """
        return self.prefix + unparse_host_port(*self.get_host_port())

    @property
    def contact_address(self):
        """
        The contact address as a string.
        """
        host, port = self.get_host_port()
        host = ensure_concrete_host(host)
        return self.prefix + unparse_host_port(host, port)


class TCPListener(BaseTCPListener):
    prefix = 'tcp://'
    comm_class = TCP
    encrypted = False

    def _get_server_args(self, **connection_args):
        return {}

    @gen.coroutine
    def _prepare_stream(self, stream, address):
        raise gen.Return(stream)


class TLSListener(BaseTCPListener):
    prefix = 'tls://'
    comm_class = TLS
    encrypted = True

    def _get_server_args(self, **connection_args):
        ctx = _expect_tls_context(connection_args)
        return {'ssl_options': ctx}

    @gen.coroutine
    def _prepare_stream(self, stream, address):
        try:
            yield stream.wait_for_handshake()
        except EnvironmentError as e:
            # The handshake went wrong, log and ignore
            logger.warning("Listener on %r: TLS handshake failed with remote %r: %s",
                           self.listen_address, address,
                           getattr(e, "real_error", None) or e)
        else:
            raise gen.Return(stream)


class BaseTCPBackend(Backend):

    # I/O

    def get_connector(self):
        return self._connector_class()

    def get_listener(self, loc, handle_comm, deserialize, **connection_args):
        return self._listener_class(loc, handle_comm, deserialize, **connection_args)

    # Address handling

    def get_address_host(self, loc):
        return parse_host_port(loc)[0]

    def get_address_host_port(self, loc):
        return parse_host_port(loc)

    def resolve_address(self, loc):
        host, port = parse_host_port(loc)
        return unparse_host_port(ensure_ip(host), port)

    def get_local_address_for(self, loc):
        host, port = parse_host_port(loc)
        host = ensure_ip(host)
        if ':' in host:
            local_host = get_ipv6(host)
        else:
            local_host = get_ip(host)
        return unparse_host_port(local_host, None)


class TCPBackend(BaseTCPBackend):
    _connector_class = TCPConnector
    _listener_class = TCPListener


class TLSBackend(BaseTCPBackend):
    _connector_class = TLSConnector
    _listener_class = TLSListener


backends['tcp'] = TCPBackend()
backends['tls'] = TLSBackend()
