# *****************************************************************************
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Markus Zolliker <markus.zolliker@psi.ch>
#
# *****************************************************************************

"""asynchronous connections

generic class for byte oriented communication
includes implementation for TCP and Serial connections
support for asynchronous communication, but may be used also for
synchronous IO (see frappy.io)
"""

import ast
import socket
import time
import re
from select import select

from frappy.errors import CommunicationFailedError, ConfigError
from frappy.lib import closeSocket, parse_host_port, SECoP_DEFAULT_PORT

try:
    from serial import Serial
except ImportError:
    Serial = None


class ConnectionClosed(ConnectionError):
    pass


class AsynConn:
    timeout = 1  # inter byte timeout
    scheme = None
    SCHEME_MAP = {}
    _connection = None  # None means disconnected

    def __new__(cls, uri, end_of_line=b'\n', default_settings=None):
        scheme = uri.split('://')[0]
        iocls = cls.SCHEME_MAP.get(scheme, None)
        if not iocls:
            # try tcp, if scheme not given
            try:
                parse_host_port(uri, 1)  # check hostname only
            except ValueError:
                if 'COM' in uri:
                    raise ValueError("the correct uri for a COM port is: "
                                     "'serial://COM<i>[?<option>=<value>[&<option>=value ...]]'") from None
                if '/dev' in uri:
                    raise ValueError("the correct uri for a serial port is: "
                                     "'serial:///dev/<tty>[?<option>=<value>[&<option>=value ...]]'") from None
                raise ValueError(f'invalid hostname {uri!r}') from None
            iocls = cls.SCHEME_MAP['tcp']
        return object.__new__(iocls)

    def __init__(self, uri, end_of_line=b'\n', default_settings=None):
        self.end_of_line = end_of_line
        self.default_settings = default_settings or {}
        self._rxbuffer = b''

    def __del__(self):
        self.disconnect()

    @classmethod
    def __init_subclass__(cls):
        """register subclass to scheme, if available"""
        if cls.scheme:
            cls.SCHEME_MAP[cls.scheme] = cls

    @property
    def connection(self):
        if self._connection is None:
            raise ConnectionClosed()
        return self._connection

    def shutdown(self):
        """prepare connection for disconnect, can be empty"""

    def disconnect(self):
        raise NotImplementedError

    def send(self, data):
        """send data (bytes!)

        tries to send all data"""
        raise NotImplementedError

    def recv(self):
        """return bytes received within timeout

        in contrast to socket.recv:
        - returns b'' on timeout
        - raises ConnectionClosed if the other end has disconnected
        """
        raise NotImplementedError

    def recv_nowait(self):
        """return bytes in buffer without waiting"""
        raise NotImplementedError

    def flush_recv(self):
        """flush all available bytes (return them)"""
        result = self._rxbuffer + self.recv_nowait()
        self._rxbuffer = b''
        return result

    def readline(self, timeout=None):
        """read one line

        Return either a complete line or None if not enough is available
        within 1 sec (self.timeout).
        With timeout=0 no waiting happens at all.
        If a non-zero timeout is given, a timeout error is raised instead
        of returning None.
        The timeout resolution is self.timeout (1 sec by default)
        """
        end = time.time() + timeout if timeout else 0
        recv = self.recv_nowait if timeout == 0 else self.recv
        while True:
            splitted = self._rxbuffer.split(self.end_of_line, 1)
            if len(splitted) == 2:
                line, self._rxbuffer = splitted
                return line
            data = recv()
            if not data:
                if timeout:
                    if time.time() < end:
                        continue
                    raise TimeoutError(f'timeout in readline ({timeout:g} sec)')
                return None
            self._rxbuffer += data

    def readbytes(self, nbytes, timeout=None):
        """read a fixed number of bytes

        Return either <nbytes> bytes or None if not enough data is available
        within 1 sec (self.timeout).
        With timeout=0 no waiting happens at all.
        If a non-zero timeout is given, a timeout error is raised instead
        of returning None
        The timeout resolution is self.timeout (1 sec by default)
        """
        end = time.time() + timeout if timeout else 0
        recv = self.recv_nowait if timeout == 0 else self.recv
        while len(self._rxbuffer) < nbytes:
            data = recv()
            if not data:
                if timeout:
                    if time.time() < end:
                        continue
                    raise TimeoutError(f'timeout in readbytes ({timeout:g} sec)')
                return None
            self._rxbuffer += data
        line = self._rxbuffer[:nbytes]
        self._rxbuffer = self._rxbuffer[nbytes:]
        return line

    def writeline(self, line):
        self.send(line + self.end_of_line)

    def read_ready(self, timeout=1):
        return bool(select([self.connection], [], [], timeout)[0])


class AsynTcp(AsynConn):
    """a tcp/ip connection

    uri syntax::

       tcp://<host address>:<port number>
    """
    scheme = 'tcp'

    def __init__(self, uri, *args, **kwargs):
        super().__init__(uri, *args, **kwargs)
        self.uri = uri
        if uri.startswith('tcp://'):
            uri = uri[6:]
        try:
            host, port = parse_host_port(
                uri, self.default_settings.get('port', SECoP_DEFAULT_PORT))
            self._connection = socket.create_connection(
                (host, port), timeout=self.timeout)
        except (ConnectionRefusedError, socket.gaierror, socket.timeout) as e:
            # indicate that retrying might make sense
            raise CommunicationFailedError(f'can not connect to {host}:{port}, {e}') from None

    def shutdown(self):
        if self._connection:
            try:
                self.connection.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass  # in case socket is already disconnected

    def disconnect(self):
        if self._connection:
            closeSocket(self.connection)
        self._connection = None

    def send(self, data):
        """send data (bytes!)"""
        # remark: will raise socket.timeout when output buffer is full and blocked for self.timeout
        self.connection.sendall(data)

    def recv(self):
        """return bytes in the recv buffer

        or bytes received within self.timeout
        """
        try:
            data = self.connection.recv(1024*1024)
            if data:
                return data
        except (socket.timeout, TimeoutError):
            # timeout while waiting
            return b''
        except ConnectionResetError:
            pass  # treat equally as a gracefully disconnected peer
        # note that when no data is sent on a connection, an interruption might
        # not be detected within a reasonable time. sending a heartbeat should
        # help in this case.
        raise ConnectionClosed()  # marks end of connection

    def recv_nowait(self):
        """return bytes in the recv buffer"""
        if select([self.connection], [], [], 0)[0]:
            return self.recv()
        return b''


class AsynSerial(AsynConn):
    """a serial connection using pyserial

    uri syntax::

      serial://<serial device>?[<option>=<value>[&<option>=<value> ...]]

      options (defaults, other examples):

      baudrate=9600  # 4800, 115200
      bytesize=8     # 5,6,7
      parity=none    # even, odd, mark, space
      stopbits=1     # 1.5, 2
      xonxoff=False  # True

      and others (see documentation of serial.Serial)
    """
    scheme = 'serial'
    PARITY_NAMES = {name[0]: name for name in ['NONE', 'ODD', 'EVEN', 'MARK', 'SPACE']}
    ARG_SEP = re.compile('[+&]')  # allow + or & as options separator in uri
    SETTINGS = set(Serial(None).get_settings()) if Serial else None  # keys of valid Serial settings

    def __init__(self, uri, *args, **kwargs):
        if Serial is None:
            raise ConfigError('pyserial is not installed')
        super().__init__(uri, *args, **kwargs)
        self.uri = uri
        if uri.startswith('serial://'):
            # should be the case always
            uri = uri[9:]
        uri = uri.split('?', 1)
        dev = uri[0]
        options = {k: v for k, v in self.default_settings.items() if k in self.SETTINGS}
        if len(uri) > 1:
            for kv in self.ARG_SEP.split(uri[1]):
                try:
                    key, value = kv.split('=')
                except TypeError:
                    raise ConfigError(f'{kv!r} must be <key>=<value>') from None
                if key == 'parity':
                    options[key] = value
                else:
                    options[key] = ast.literal_eval(value.title())  # title(): turn false/true into False/True
        parity = options.get('parity')
        if parity:
            name = parity.upper()
            fullname = self.PARITY_NAMES[name[0]]
            if not fullname.startswith(name):
                raise ConfigError(f'illegal parity: {parity}')
            options['parity'] = name[0]
        if 'timeout' not in options:
            options['timeout'] = self.timeout
        try:
            self._connection = Serial(dev, **options)
        except ValueError as e:
            raise ConfigError(e) from None
        # TODO: turn exceptions into ConnectionFailedError, where a retry makes sense

    def disconnect(self):
        if self._connection:
            self.connection.close()
        self._connection = None

    def send(self, data):
        """send data (bytes!)"""
        self.connection.write(data)

    def recv(self):
        """return bytes received within self.timeout"""
        n = self.connection.in_waiting
        if n:
            return self.connection.read(n)
        data = self.connection.read(1)
        return data + self.connection.read(self.connection.in_waiting)

    def recv_nowait(self):
        n = self.connection.in_waiting
        return self.connection.read(n) if n else b''

    if not hasattr(Serial, 'fileno'):
        # select is not supported for serial objects (Windows)
        # create a less efficient workaround
        def read_ready(self, timeout=1):
            if self.connection.in_waiting:
                return True
            deadline = time.time() + timeout
            while time.time() < deadline:
                if self.connection.in_waiting:
                    return True
                time.sleep(0.05)
            return False
