import collections
import dataclasses
import os
import queue
import selectors
import socket

from . import buffers
from . import contextlib
from . import thread


class Connection:
    def __init__(self, sock, addr):
        self.sock = sock
        self.addr = addr
        self.recvbuf = None
        self.handshake = b''
        self.sendbufs = collections.deque()

    def fileno(self):
        return self.sock.fileno()


@dataclasses.dataclass
class Options:
    ipv6: bool = False
    host: str = ''
    max_msg_size: int = 4 * 1024**3
    max_recv_queue: int = 4096
    max_send_queue: int = 4096
    logging: bool = True
    logging_color: str = 'blue'
    handshake: str = 'portal_handshake'


class ServerSocket:
    def __init__(self, port, name='Server', **kwargs):
        if isinstance(port, str):
            assert '://' not in port, port
            port = int(port.rsplit(':', 1)[-1])
        self.name = name
        self.options = Options(**{**contextlib.context.serverkw, **kwargs})
        self.handshake = self.options.handshake.encode('utf-8')
        if self.options.ipv6:
            self.sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            self.sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
            self.addr = (self.options.host or '::', port, 0, 0)
        else:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.addr = (self.options.host or '0.0.0.0', port)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)  # TODO
        self._log(f'Binding to {self.addr[0]}:{self.addr[1]}')
        self.sock.bind(self.addr)
        self.sock.setblocking(False)
        self.sock.listen(8192)
        self.get_signal, self.set_signal = os.pipe()
        self.sel = selectors.DefaultSelector()
        self.sel.register(self.get_signal, selectors.EVENT_READ, data='signal')
        self.sel.register(self.sock, selectors.EVENT_READ, data=None)
        self._log(f'Listening at {self.addr[0]}:{self.addr[1]}')
        self.conns = {}
        self.recvq = queue.Queue()  # [(addr, bytes)]
        self.reading = True
        self.running = True
        self.error = None
        self.thread = thread.Thread(self._loop, name=f'{name}Loop', start=True)

    @property
    def connections(self):
        return tuple(self.conns.keys())

    def recv(self, timeout=None):
        if self.error:
            raise self.error
        assert self.running
        try:
            return self.recvq.get(block=(timeout != 0), timeout=timeout)
        except queue.Empty:
            raise TimeoutError

    def send(self, addr, *data):
        if self.error:
            raise self.error
        assert self.running
        if self._numsending() > self.options.max_send_queue:
            raise RuntimeError('Too many outgoing messages enqueued')
        maxsize = self.options.max_msg_size
        try:
            buf = buffers.SendBuffer(*data, maxsize=maxsize)
            self.conns[addr].sendbufs.append(buf)
            os.write(self.set_signal, bytes(1))
        except KeyError:
            self._log('Dropping message to disconnected client')

    def shutdown(self):
        self.reading = False

    def close(self, timeout=None):
        self.running = False
        self.thread.join(timeout)
        [conn.sock.close() for conn in self.conns.values()]
        self.sock.close()
        self.sel.close()
        os.close(self.get_signal)
        os.close(self.set_signal)

    def _loop(self):
        writing = False
        try:
            while self.running or self._numsending():
                for key, mask in self.sel.select(timeout=0.2):
                    if key.data == 'signal':
                        writing = True
                        os.read(self.get_signal, 1)
                    elif key.data is None and self.reading:
                        assert mask & selectors.EVENT_READ
                        self._accept(key.fileobj)
                    elif mask & selectors.EVENT_READ and self.reading:
                        self._recv(key.data)
                if not writing:
                    continue
                pending = [
                    conn for conn in self.conns.values() if conn.sendbufs
                ]
                for conn in pending:
                    try:
                        conn.sendbufs[0].send(conn.sock)
                        if conn.sendbufs[0].done():
                            conn.sendbufs.popleft()
                            if not any(conn.sendbufs for conn in pending):
                                writing = False
                    except BlockingIOError:
                        pass
                    except ConnectionResetError:
                        # The client is gone but we may have buffered messages left
                        # to read, so we keep the socket open until recv() fails.
                        pass
        except Exception as e:
            self.error = e

    def _accept(self, sock):
        sock, addr = sock.accept()
        self._log(f'Accepted connection from {addr[0]}:{addr[1]}')
        sock.setblocking(False)
        conn = Connection(sock, addr)
        self.sel.register(sock, selectors.EVENT_READ, data=conn)
        self.conns[addr] = conn

    def _recv(self, conn):
        if not conn.recvbuf:
            conn.recvbuf = buffers.RecvBuffer(self.options.max_msg_size)
        try:
            if len(conn.handshake) < len(self.handshake):
                self._handshake(conn)
                return
            else:
                conn.recvbuf.recv(conn.sock)
        except OSError as e:
            # For example:
            # - ConnectionResetError
            # - TimeoutError: [Errno 110] Connection timed out
            self._disconnect(conn, e)
            return
        if not conn.recvbuf.done():
            return
        if self.recvq.qsize() > self.options.max_recv_queue:
            raise RuntimeError('Too many incoming messages enqueued')
        self.recvq.put((conn.addr, conn.recvbuf.result()))
        conn.recvbuf = None

    def _disconnect(self, conn, e):
        detail = f'{type(e).__name__}'
        detail = f'{detail}: {e}' if str(e) else detail
        self._log(
            f'Closed connection to {conn.addr[0]}:{conn.addr[1]} ({detail})'
        )
        conn = self.conns.pop(conn.addr)
        if conn.sendbufs:
            count = len(conn.sendbufs)
            conn.sendbufs.clear()
            self._log(f'Dropping {count} messages to disconnected client')
        self.sel.unregister(conn.sock)
        conn.sock.close()

    def _numsending(self):
        return sum(len(x.sendbufs) for x in self.conns.values())

    def _handshake(self, conn):
        assert len(conn.handshake) < len(self.handshake)
        size = len(self.handshake) - len(conn.handshake)
        conn.handshake += conn.sock.recv(size)
        if conn.handshake != self.handshake[: len(conn.handshake)]:
            m = f"Expected handshake '{self.handshake}' got '{conn.handshake}'"
            self._disconnect(conn, ValueError(m))

    def _log(self, *args, **kwargs):
        if not self.options.logging:
            return
        color = self.options.logging_color
        contextlib.context.print(self.name, *args, color=color)
