#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals


from .TTransport import TTransportBase, TTransportException, \
        TServerTransportBase
import os
import sys
import errno
import select
import socket
import warnings
import time

try:
    import fcntl
except ImportError:
    # Windows doesn't have this module
    fcntl = None

class ConnectionEpoll:
    """ epoll is preferred over select due to its efficiency and ability to
        handle more than 1024 simultaneous connections """

    def __init__(self):
        self.epoll = select.epoll()
        # TODO should we set any other masks?
        # http://docs.python.org/library/select.html#epoll-objects
        self.READ_MASK = select.EPOLLIN | select.EPOLLPRI
        self.WRITE_MASK = select.EPOLLOUT
        self.ERR_MASK = select.EPOLLERR | select.EPOLLHUP

    def read(self, fileno):
        self.unregister(fileno)
        self.epoll.register(
            fileno, self.READ_MASK | self.ERR_MASK
        )

    def write(self, fileno):
        self.unregister(fileno)
        self.epoll.register(fileno, self.WRITE_MASK)

    def unregister(self, fileno):
        try:
            self.epoll.unregister(fileno)
        except:
            pass

    def process(self, timeout):
        # poll() invokes a "long" syscall that will be interrupted by any signal
        # that comes in, causing an EINTR error.  If this happens, avoid dying
        # horribly by trying again with the appropriately shortened timout.
        deadline = time.clock() + float(timeout or 0)
        poll_timeout = float(timeout or -1)
        while True:
            if timeout is not None and timeout > 0:
                poll_timeout = max(0, deadline - time.clock())
            try:
                msgs = self.epoll.poll(timeout=poll_timeout)
                break
            except IOError as e:
                if e.errno == errno.EINTR:
                    continue
                else:
                    raise

        rset = []
        wset = []
        xset = []
        for fd, mask in msgs:
            if mask & self.READ_MASK:
                rset.append(fd)
            if mask & self.WRITE_MASK:
                wset.append(fd)
            if mask & self.ERR_MASK:
                xset.append(fd)

        return rset, wset, xset


class ConnectionSelect:
    def __init__(self):
        self.readable = set()
        self.writable = set()

    def read(self, fileno):
        if fileno in self.writable:
            self.writable.remove(fileno)
        self.readable.add(fileno)

    def write(self, fileno):
        if fileno in self.readable:
            self.readable.remove(fileno)
        self.writable.add(fileno)

    def unregister(self, fileno):
        if fileno in self.readable:
            self.readable.remove(fileno)
        elif fileno in self.writable:
            self.writable.remove(fileno)

    def registered(self, fileno):
        return fileno in self.readable or fileno in self.writable

    def process(self, timeout):
        # select() invokes a "long" syscall that will be interrupted by any
        # signal that comes in, causing an EINTR error.  If this happens,
        # avoid dying horribly by trying again with the appropriately
        # shortened timout.
        deadline = time.clock() + float(timeout or 0)
        poll_timeout = timeout if timeout is None or timeout > 0 else None
        while True:
            if timeout is not None and timeout > 0:
                poll_timeout = max(0, deadline - time.clock())
            try:
                return select.select(list(self.readable), list(self.writable),
                        list(self.readable), poll_timeout)
            except IOError as e:
                if e.errno == errno.EINTR:
                    continue
                else:
                    raise


class TSocketBase(TTransportBase):
    """Base class for both connected and listening sockets"""
    def __init__(self):
        self.handles = {}

    def _resolveAddr(self, family=None):
        if family is None:
            family = socket.AF_UNSPEC
        if self._unix_socket is not None:
            return [(socket.AF_UNIX, socket.SOCK_STREAM, None, None,
                     self._unix_socket)]
        else:
            ai_flags = 0
            if self.host is None:
                ai_flags |= socket.AI_PASSIVE
            return socket.getaddrinfo(self.host, self.port, family,
                                      socket.SOCK_STREAM, 0,
                                      ai_flags)

    def close(self):
        klist = self.handles.keys() if sys.version_info[0] < 3 else \
            list(self.handles.keys())
        for key in klist:
            self.handles[key].close()
            del self.handles[key]

    def getSocketName(self):
        if not self.handles:
            raise TTransportException(TTransportException.NOT_OPEN,
                'Transport not open')
        return next(iter(self.handles.values())).getsockname()

    def fileno(self):
        if not self.handles:
            raise TTransportException(TTransportException.NOT_OPEN,
                'Transport not open')
        if sys.version_info[0] >= 3:
            return list(self.handles.values())[0].fileno()
        else:
            return self.handles.values()[0].fileno()

    def setCloseOnExec(self, closeOnExec):
        self.close_on_exec = closeOnExec
        for handle in self.handles.values():
            self._setHandleCloseOnExec(handle)

    def _setHandleCloseOnExec(self, handle):
        # Windows doesn't have this module, don't set the handle in this case.
        if fcntl is None: