"""
Low level socket utilities.

Pyro - Python Remote Objects.  Copyright by Irmen de Jong (irmen@razorvine.net).
"""

import os
import platform
import socket
import errno
import time
import select
import ipaddress
import weakref
import contextlib
from typing import Union, Optional, Tuple, Dict, Type, Any
try:
    import ssl
except ImportError:
    pass
from . import config
from .errors import CommunicationError, TimeoutError, ConnectionClosedError


# Note: other interesting errnos are EPERM, ENOBUFS, EMFILE
# but it seems to me that all these signify an unrecoverable situation.
# So I didn't include them in the list of retryable errors.
ERRNO_RETRIES = [errno.EINTR, errno.EAGAIN, errno.EWOULDBLOCK, errno.EINPROGRESS]
if hasattr(errno, "WSAEINTR"):
    ERRNO_RETRIES.append(errno.WSAEINTR)
if hasattr(errno, "WSAEWOULDBLOCK"):
    ERRNO_RETRIES.append(errno.WSAEWOULDBLOCK)
if hasattr(errno, "WSAEINPROGRESS"):
    ERRNO_RETRIES.append(errno.WSAEINPROGRESS)

ERRNO_BADF = [errno.EBADF]
if hasattr(errno, "WSAEBADF"):
    ERRNO_BADF.append(errno.WSAEBADF)

ERRNO_ENOTSOCK = [errno.ENOTSOCK]
if hasattr(errno, "WSAENOTSOCK"):
    ERRNO_ENOTSOCK.append(errno.WSAENOTSOCK)
if not hasattr(socket, "SOL_TCP"):
    socket.SOL_TCP = socket.IPPROTO_TCP

ERRNO_EADDRNOTAVAIL = [errno.EADDRNOTAVAIL]
if hasattr(errno, "WSAEADDRNOTAVAIL"):
    ERRNO_EADDRNOTAVAIL.append(errno.WSAEADDRNOTAVAIL)

ERRNO_EADDRINUSE = [errno.EADDRINUSE]
if hasattr(errno, "WSAEADDRINUSE"):
    ERRNO_EADDRINUSE.append(errno.WSAEADDRINUSE)

# msg_waitall has proven to be unreliable on windows
USE_MSG_WAITALL = hasattr(socket, "MSG_WAITALL") and platform.system() != "Windows"


def get_ip_address(hostname: str, workaround127: bool = False, version: int = None) \
        -> Union[ipaddress.IPv4Address, ipaddress.IPv6Address]:
    """
    Returns the IP address for the given host. If you enable the workaround,
    it will use a little hack if the ip address is found to be the loopback address.
    The hack tries to discover an externally visible ip address instead (this only works for ipv4 addresses).
    Set ipVersion=6 to return ipv6 addresses, 4 to return ipv4, 0 to let OS choose the best one or None to use config.PREFER_IP_VERSION.
    """
    if not workaround127:
        with contextlib.suppress(ValueError):
            addr = ipaddress.ip_address(hostname)
            return addr

    def getaddr(ip_version):
        if ip_version == 6:
            family = socket.AF_INET6
        elif ip_version == 4:
            family = socket.AF_INET
        elif ip_version == 0:
            family = socket.AF_UNSPEC
        else:
            raise ValueError("unknown value for argument ipVersion.")
        ip = socket.getaddrinfo(hostname or socket.gethostname(), 80, family, socket.SOCK_STREAM, socket.SOL_TCP)[0][4][0]
        if workaround127 and (ip.startswith("127.") or ip == "0.0.0.0"):
            return get_interface("4.2.2.2").ip
        return ipaddress.ip_address(ip)

    try:
        if hostname and ':' in hostname and version is None:
            version = 0
        return getaddr(config.PREFER_IP_VERSION) if version is None else getaddr(version)
    except socket.gaierror:
        if version == 6 or (version is None and config.PREFER_IP_VERSION == 6):
            raise socket.error("unable to determine IPV6 address")
        return getaddr(0)


def get_interface(ip_address: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address]) \
        -> Union[ipaddress.IPv4Interface, ipaddress.IPv6Interface]:
    """tries to find the network interface that connects to the given host's address"""
    if isinstance(ip_address, str):
        ip_address = get_ip_address(ip_address)
    family = socket.AF_INET if ip_address.version == 4 else socket.AF_INET6
    with socket.socket(family, socket.SOCK_DGRAM) as sock:
        sock.connect((str(ip_address), 53))  # 53=dns
        return ipaddress.ip_interface(sock.getsockname()[0])


def __retrydelays():
    # first try a few very short delays,
    # if that doesn't work, increase by 0.1 sec every time
    yield 0.0001
    yield 0.001
    yield 0.01
    d = 0.1
    while True:
        yield d
        d += 0.1


def receive_data(sock: socket.socket, size: int) -> bytes:
    """Retrieve a given number of bytes from a socket.
    It is expected the socket is able to supply that number of bytes.
    If it isn't, an exception is raised (you will not get a zero length result
    or a result that is smaller than what you asked for). The partial data that
    has been received however is stored in the 'partialData' attribute of
    the exception object."""
    try:
        delays = __retrydelays()
        msglen = 0
        data = bytearray()
        if USE_MSG_WAITALL and not hasattr(sock, "getpeercert"):    # ssl doesn't support recv flags
            while True:
                try:
                    chunk = sock.recv(size, socket.MSG_WAITALL)
                    if len(chunk) == size:
                        return chunk
                    # less data than asked, drop down into normal receive loop to finish
                    msglen = len(chunk)
                    data.extend(chunk)
                    break
                except socket.timeout:
                    raise TimeoutError("receiving: timeout")
                except socket.error as x:
                    err = getattr(x, "errno", x.args[0])
                    if err not in ERRNO_RETRIES:
                        raise ConnectionClosedError("receiving: connection lost: " + str(x))
                    time.sleep(next(delays))  # a slight delay to wait before retrying
        # old fashioned recv loop, we gather chunks until the message is complete
        while True:
            try:
                while msglen < size:
                    # 60k buffer limit avoids problems on certain OSes like VMS, Windows
                    chunk = sock.recv(min(60000, size - msglen))
                    if not chunk:
                        break
                    data.extend(chunk)
                    msglen += len(chunk)
                if len(data) != size:
                    err = ConnectionClosedError("receiving: not enough data")
                    err.partialData = data  # store the message that was received until now
                    raise err
                return data  # yay, complete
            except socket.timeout:
                raise TimeoutError("receiving: timeout")
            except socket.error as x:
                err = getattr(x, "errno", x.args[0])
                if err not in ERRNO_RETRIES:
                    raise ConnectionClosedError("receiving: connection lost: " + str(x))
                time.sleep(next(delays))  # a slight delay to wait before retrying
    except socket.timeout:
        raise TimeoutError("receiving: timeout")


def send_data(sock: socket.socket, data: bytes) -> None:
    """
    Send some data over a socket.
    Some systems have problems with ``sendall()`` when the socket is in non-blocking mode.
    For instance, Mac OS X seems to be happy to throw EAGAIN errors too often.
    This function falls back to using a regular send loop if needed.
    """
    if sock.gettimeout() is None:
        # socket is in blocking mode, we can use sendall normally.
        try:
            sock.sendall(data)
            return
        except socket.timeout:
            raise TimeoutError("sending: timeout")
        except socket.error as x:
            raise ConnectionClosedError("sending: connection lost: " + str(x))
    else:
        # Socket is in non-blocking mode, use regular send loop.
        delays = __retrydelays()
        while data:
            try:
                sent = sock.send(data)
                data = data[sent:]
            except socket.timeout:
                raise TimeoutError("sending: timeout")
            except socket.error as x:
                err = getattr(x, "errno", x.args[0])
                if err not in ERRNO_RETRIES:
                    raise ConnectionClosedError("sending: connection lost: " + str(x))
                time.sleep(next(delays))  # a slight delay to wait before retrying


def create_socket(bind: Union[Tuple, str] = None,
                  connect: Union[Tuple, str] = None,
                  reuseaddr: bool = False, keepalive: bool = True,
                  timeout: Optional[float] = -1, noinherit: bool = False,
                  ipv6: bool = False, nodelay: bool = True, sslContext: ssl.SSLContext = None) -> socket.socket:
    """
    Create a socket. Default socket options are keepalive and IPv4 family, and nodelay (nagle disabled).
    If 'bind' or 'connect' is a string, it is assumed a Unix domain socket is requested.
    Otherwise, a normal tcp/ip socket tuple (addr, port, ...) is used.
    Set ipv6=True to create an IPv6 socket rather than IPv4.
    Set ipv6=None to use the PREFER_IP_VERSION config setting.
    """
    if bind and connect:
        raise ValueError("bind and connect cannot both be specified at the same time")
    forceIPv6 = ipv6 or (ipv6 is None and config.PREFER_IP_VERSION == 6)
    if isinstance(bind, str) or isinstance(connect, str):
        family = socket.AF_UNIX
    elif not bind and not connect:
        family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
    elif isinstance(bind, tuple):
        if not bind[0]:
            family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
        else:
            addr = get_ip_address(bind[0])
            if addr.version == 4:
                if forceIPv6:
                    raise ValueError("IPv4 address is used bind argument with forceIPv6 argument:" + bind[0] + ".")
                family = socket.AF_INET
            elif addr.version == 6:
                family = socket.AF_INET6
                # replace bind addresses by their ipv6 counterparts (4-tuple)
                bind = (bind[0], bind[1], 0, 0)
            else:
                raise ValueError("unknown bind format.")
    elif isinstance(connect, tuple):
        if not connect[0]:
            family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
        else:
            addr = get_ip_address(connect[0])
            if addr.version == 4:
                if forceIPv6:
                    raise ValueError("IPv4 address is used in connect argument with forceIPv6 argument:" + connect[0] + ".")
                family = socket.AF_INET
            elif addr.version == 6:
                family = socket.AF_INET6
                # replace connect addresses by their ipv6 counterparts (4-tuple)
                connect = (connect[0], connect[1], 0, 0)
            else:
                raise ValueError("unknown connect format.")
    else:
        raise ValueError("unknown bind or connect format.")
    sock = socket.socket(family, socket.SOCK_STREAM)
    if sslContext:
        if bind:
            sock = sslContext.wrap_socket(sock, server_side=True)
        elif connect:
            sock = sslContext.wrap_socket(sock, server_side=False, server_hostname=connect[0])
        else:
            sock = sslContext.wrap_socket(sock, server_side=False)
    if nodelay:
        set_nodelay(sock)
    if reuseaddr:
        set_reuseaddr(sock)
    if noinherit:
        set_noinherit(sock)
    if timeout is not None:
        if timeout == 0:
            timeout = None
        elif timeout >= 0:
            sock.settimeout(timeout)
    if bind:
        if type(bind) is tuple and bind[1] == 0:
            bind_unused_port(sock, bind[0])
        else:
            sock.bind(bind)
        with contextlib.suppress(OSError, IOError):
            sock.listen(100)
    if connect:
        try:
            sock.connect(connect)
        except socket.error as xv:
            # This can happen when the socket is in non-blocking mode (or has a timeout configured).
            # We check if it is a retryable errno (usually EINPROGRESS).
            # If so, we use select() to wait until the socket is in writable state,
            # essentially rebuilding a blocking connect() call.
            errno = getattr(xv, "errno", 0)
            if errno in ERRNO_RETRIES:
                if timeout:
                    if timeout < 0:
                        timeout = None
                    elif timeout < 0.1:
                        timeout = 0.1
                while True:
                    try:
                        sr, sw, se = select.select([], [sock], [sock], timeout)
                    except InterruptedError:
                        continue
                    if sock in sw:
                        break  # yay, writable now, connect() completed
                    elif sock in se:
                        sock.close()  # close the socket that refused to connect
                        raise socket.error("connect failed")
            else:
                sock.close()  # close the socket that refused to connect
                raise
    if keepalive:
        set_keepalive(sock)
    return sock


def create_bc_socket(bind: Union[Tuple, str] = None, reuseaddr: bool = False,
                     timeout: Optional[float] = -1, ipv6: bool = False) -> socket.socket:
    """
    Create a udp broadcast socket.
    Set ipv6=True to create an IPv6 socket rather than IPv4.
    Set ipv6=None to use the PREFER_IP_VERSION config setting.
    """
    forceIPv6 = ipv6 or (ipv6 is None and config.PREFER_IP_VERSION == 6)
    if not bind:
        family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
    elif isinstance(bind, tuple):
        if not bind[0]:
            family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
        else:
            addr = get_ip_address(bind[0])
            if addr.version == 4:
                if forceIPv6:
                    raise ValueError("IPv4 address is used with forceIPv6 option:" + bind[0] + ".")
                family = socket.AF_INET
            elif addr.version == 6:
                family = socket.AF_INET6
                bind = (bind[0], bind[1], 0, 0)
            else:
                raise ValueError("unknown bind format: %r" % (bind,))
    else:
        raise ValueError("unknown bind format: %r" % (bind,))
    sock = socket.socket(family, socket.SOCK_DGRAM)
    if family == socket.AF_INET:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    if reuseaddr:
        set_reuseaddr(sock)
    if timeout is None:
        sock.settimeout(None)
    else:
        if timeout >= 0:
            sock.settimeout(timeout)
    if bind:
        host = bind[0] or ""
        port = bind[1]
        if port == 0:
            bind_unused_port(sock, host)
        else:
            if len(bind) == 2:
                sock.bind((host, port))  # ipv4
            elif len(bind) == 4:
                sock.bind((host, port, 0, 0))  # ipv6
            else:
                raise ValueError("bind must be None, 2-tuple or 4-tuple")
    return sock


def set_reuseaddr(sock: socket.socket) -> None:
    """sets the SO_REUSEADDR option on the socket, if possible."""
    with contextlib.suppress(Exception):
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)


def set_nodelay(sock: socket.socket) -> None:
    """sets the TCP_NODELAY option on the socket (to disable Nagle's algorithm), if possible."""
    with contextlib.suppress(Exception):
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)


def set_keepalive(sock: socket.socket) -> None:
    """sets the SO_KEEPALIVE option on the socket, if possible."""
    with contextlib.suppress(Exception):
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)


try:
    import fcntl

    def set_noinherit(sock: socket.socket) -> None:
        """Mark the given socket fd as non-inheritable to child processes"""
        fd = sock.fileno()
        flags = fcntl.fcntl(fd, fcntl.F_GETFD)
        fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)

except ImportError:
    # no fcntl available, try the windows version
    try:
        from ctypes import windll, WinError, wintypes
        # help ctypes to set the proper args for this kernel32 call on 64-bit pythons
        _SetHandleInformation = windll.kernel32.SetHandleInformation
        _SetHandleInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD]
        _SetHandleInformation.restype = wintypes.BOOL  # don't need this, but might as well

        def set_noinherit(sock: socket.socket) -> None:
            """Mark the given socket fd as non-inheritable to child processes"""
            if not _SetHandleInformation(sock.fileno(), 1, 0):
                raise WinError()

    except (ImportError, NotImplementedError):
        # nothing available, define a dummy function
        def set_noinherit(sock: socket.socket) -> None:
            """Mark the given socket fd as non-inheritable to child processes (dummy)"""
            pass


class SocketConnection(object):
    """A wrapper class for plain sockets, containing various methods such as :meth:`send` and :meth:`recv`"""
    def __init__(self, sock: socket.socket, objectId: str = None, keep_open: bool = False) -> None:
        self.sock = sock
        self.objectId = objectId
        self.pyroInstances = {}    # type: Dict[Type, Any]   # pyro objects for instance_mode=session
        self.tracked_resources = weakref.WeakSet()   # type: weakref.WeakSet[Any]  # weakrefs to resources for this connection
        self.keep_open = keep_open

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

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

    def send(self, data: bytes) -> None:
        send_data(self.sock, data)

    def recv(self, size: int) -> bytes:
        return receive_data(self.sock, size)

    def close(self) -> None:
        if self.keep_open:
            return
        with contextlib.suppress(Exception):
            self.sock.shutdown(socket.SHUT_RDWR)
        with contextlib.suppress(Exception):
            self.sock.close()
        self.pyroInstances = {}   # release the session instances
        for rsc in self.tracked_resources:
            with contextlib.suppress(Exception):
                rsc.close()     # it is assumed a 'resource' has a close method.
        self.tracked_resources.clear()

    def fileno(self) -> int:
        return self.sock.fileno()

    def family(self) -> str:
        return family_str(self.sock)

    def settimeout(self, timeout: Optional[float]) -> None:
        self.sock.settimeout(timeout)

    def gettimeout(self) -> Optional[float]:
        return self.sock.gettimeout()

    def getpeercert(self) -> Optional[dict]:
        try:
            return self.sock.getpeercert()      # type: ignore
        except AttributeError:
            return None

    timeout = property(gettimeout, settimeout)


def family_str(sock) -> str:
    f = sock.family
    if f == socket.AF_INET:
        return "IPv4"
    if f == socket.AF_INET6:
        return "IPv6"
    if hasattr(socket, "AF_UNIX") and f == socket.AF_UNIX:
        return "Unix"
    return "???"


def find_probably_unused_port(family: int = socket.AF_INET, socktype: int = socket.SOCK_STREAM) -> int:
    """Returns an unused port that should be suitable for binding (likely, but not guaranteed).
    This code is copied from the stdlib's test.test_support module."""
    with socket.socket(family, socktype) as sock:
        return bind_unused_port(sock)


def bind_unused_port(sock: socket.socket, host: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address] = 'localhost') -> int:
    """Bind the socket to a free port and return the port number.
    This code is based on the code in the stdlib's test.test_support module."""
    if sock.family in (socket.AF_INET, socket.AF_INET6) and sock.type == socket.SOCK_STREAM:
        if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
            with contextlib.suppress(socket.error):
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
    if not isinstance(host, str):
        host = str(host)
    if sock.family == socket.AF_INET:
        if host == 'localhost':
            sock.bind(('127.0.0.1', 0))
        else:
            sock.bind((host, 0))
    elif sock.family == socket.AF_INET6:
        if host == 'localhost':
            sock.bind(('::1', 0, 0, 0))
        else:
            sock.bind((host, 0, 0, 0))
    else:
        raise CommunicationError("unsupported socket family: " + str(sock.family))
    return sock.getsockname()[1]


def interrupt_socket(address: Tuple[str, int]) -> None:
    """bit of a hack to trigger a blocking server to get out of the loop, useful at clean shutdowns"""
    with contextlib.suppress(socket.error):
        sock = create_socket(connect=address, keepalive=False, timeout=-1)
        with contextlib.suppress(socket.error, AttributeError):
            sock.sendall(b"!" * 16)
        with contextlib.suppress(OSError, socket.error):
            sock.shutdown(socket.SHUT_RDWR)
        sock.close()


__ssl_server_context = None
__ssl_client_context = None


def get_ssl_context(servercert: str = "", serverkey: str = "", clientcert: str = "", clientkey: str = "",
                    cacerts: str = "", keypassword: str = "") -> ssl.SSLContext:
    """creates an SSL context and caches it, so you have to set the parameters correctly before doing anything"""
    global __ssl_client_context, __ssl_server_context
    if servercert:
        if clientcert:
            raise ValueError("can't have both server cert and client cert")
        # server context
        if __ssl_server_context:
            return __ssl_server_context
        if not os.path.isfile(servercert):
            raise IOError("server cert file not found")
        if serverkey and not os.path.isfile(serverkey):
            raise IOError("server key file not found")
        __ssl_server_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        __ssl_server_context.load_cert_chain(servercert, serverkey or None, keypassword or None)    # type: ignore
        if cacerts:
            if os.path.isdir(cacerts):
                __ssl_server_context.load_verify_locations(capath=cacerts)
            else:
                __ssl_server_context.load_verify_locations(cafile=cacerts)
        if config.SSL_REQUIRECLIENTCERT:
            __ssl_server_context.verify_mode = ssl.CERT_REQUIRED   # 2-way ssl, server+client certs
        else:
            __ssl_server_context.verify_mode = ssl.CERT_NONE   # 1-way ssl, server cert only
        return __ssl_server_context
    else:
        # client context
        if __ssl_client_context:
            return __ssl_client_context
        __ssl_client_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        if clientcert:
            if not os.path.isfile(clientcert):
                raise IOError("client cert file not found")
            __ssl_client_context.load_cert_chain(clientcert, clientkey or None, keypassword or None)    # type: ignore
        if cacerts:
            if os.path.isdir(cacerts):
                __ssl_client_context.load_verify_locations(capath=cacerts)
            else:
                __ssl_client_context.load_verify_locations(cafile=cacerts)
        return __ssl_client_context