import contextlib
import logging
import multiprocessing
import os
import pathlib
import re
import select
import socket
import ssl as _ssl
import struct
from typing import Generator, Iterable, Optional, Type, Sequence, Dict, Tuple

import irc.bot
from irc.client import (
    Connection,
    Event as _IRCEvent,
    ServerConnection,
    DCCConnection,
    ip_numstr_to_quad,
    ip_quad_to_numstr,
)
from irc.connection import Factory as ConnectionFactory
from irc.events import codes as irc_event_codes

from platypush.context import get_bus
from platypush.message.event.irc import (
    IRCEvent,
    IRCChannelJoinEvent,
    IRCChannelKickEvent,
    IRCDisconnectEvent,
    IRCModeEvent,
    IRCNickChangeEvent,
    IRCPartEvent,
    IRCQuitEvent,
    IRCConnectEvent,
    IRCPrivateMessageEvent,
    IRCPublicMessageEvent,
    IRCDCCRequestEvent,
    IRCDCCMessageEvent,
    IRCDCCFileRequestEvent,
    IRCCTCPMessageEvent,
    IRCDCCFileRecvCompletedEvent,
    IRCDCCFileRecvCancelledEvent,
    IRCDCCFileSendCompletedEvent,
    IRCDCCFileSendCancelledEvent,
)


class IRCBot(irc.bot.SingleServerIRCBot):
    def __init__(
        self,
        server: str,
        port: int,
        nickname: str,
        alias: str,
        channels: Iterable[str],
        response_timeout: Optional[float],
        dcc_file_transfer_timeout: Optional[float],
        dcc_accept_timeout: Optional[float],
        dcc_downloads_dir: str,
        realname: Optional[str] = None,
        password: Optional[str] = None,
        ssl: bool = False,
        ipv6: bool = False,
        stop_message: Optional[str] = None,
        dcc_ip_whitelist: Optional[Sequence[str]] = None,
        dcc_ip_blacklist: Optional[Sequence[str]] = None,
        dcc_nick_whitelist: Optional[Sequence[str]] = None,
        dcc_nick_blacklist: Optional[Sequence[str]] = None,
        dcc_max_connections: Optional[int] = None,
    ):
        connection_factory = ConnectionFactory()
        if ssl:
            connection_factory.wrapper = _ssl.wrap_socket
        if ipv6:
            connection_factory.family = socket.AF_INET6

        super().__init__(
            server_list=[(server, port)],
            nickname=nickname,
            realname=realname,
            connect_factory=connection_factory,
        )

        self.server = server
        self.port = port
        self.alias = alias
        self._password = password
        self.channels.update({channel: None for channel in channels})
        self._stop_message = stop_message
        self.dcc_ip_whitelist = set(dcc_ip_whitelist or [])
        self.dcc_ip_blacklist = set(dcc_ip_blacklist or [])
        self.dcc_nick_whitelist = set(dcc_nick_whitelist or [])
        self.dcc_nick_blacklist = set(dcc_nick_blacklist or [])
        self.dcc_downloads_dir = dcc_downloads_dir
        self.response_timeout = response_timeout
        self.dcc_file_transfer_timeout = dcc_file_transfer_timeout
        self.dcc_accept_timeout = dcc_accept_timeout
        self.dcc_max_connections = dcc_max_connections
        self._dcc_send_procs: Dict[Tuple[str, int], multiprocessing.Process] = {}
        self._dcc_recv_procs: Dict[Tuple[str, int], multiprocessing.Process] = {}
        self._dcc_proc_completion_queue = multiprocessing.Queue()
        self._dcc_processor: Optional[multiprocessing.Process] = None
        self.logger = logging.getLogger(f'irc@{server}')
        # Maps <matching_event_type> -> <response_queue>
        self._pending_requests: Dict[str, multiprocessing.Queue] = {}

    @property
    def stop_message(self) -> Optional[str]:
        return self._stop_message

    def _post_event(
        self, connection: Connection, output_event_type: Type[IRCEvent], **kwargs
    ):
        if isinstance(connection, ServerConnection):
            kwargs['server'] = connection.server
            kwargs['port'] = connection.port

        kwargs['connected'] = connection.connected
        kwargs['alias'] = self.alias
        event = output_event_type(**kwargs)
        get_bus().post(event)

    def on_join(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCChannelJoinEvent,
            channel=event.target,
            nick=event.source.nick,
        )
        super()._on_join(connection, event)

    def on_kick(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCChannelKickEvent,
            channel=event.target,
            source_nick=event.arguments[1],
            target_nick=event.arguments[0],
        )

        super()._on_kick(connection, event)

    def on_nick(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCNickChangeEvent,
            before=event.source.nick,
            after=event.target.nick,
        )
        super()._on_nick(connection, event)

    def on_mode(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCModeEvent,
            mode=event.arguments[0],
            source=event.source.nick,
            target_=(event.arguments[1] if len(event.arguments) > 1 else None),
            channel=event.target,
        )

        super()._on_mode(connection, event)

    def on_part(self, connection: Connection, event: _IRCEvent):
        self._post_event(connection, IRCPartEvent, nick=event.source.nick)
        super()._on_part(connection, event)

    def on_quit(self, connection: Connection, event: _IRCEvent):
        self._post_event(connection, IRCQuitEvent, nick=event.source.nick)
        super()._on_quit(connection, event)

    def on_welcome(self, connection: Connection, *_):
        self._post_event(connection, IRCConnectEvent)
        for channel in self.channels:
            self.logger.info(f'Joining channel {channel}')
            self.connection.join(channel)

    @staticmethod
    def on_nicknameinuse(connection: ServerConnection, *_):
        connection.nick(connection.nickname + '_')

    @staticmethod
    def _mentions_me(connection: ServerConnection, text: str) -> bool:
        return bool(re.search(fr'(^|\s|@){connection.nickname}(:|\s|!|\?|$)', text))

    def on_pubmsg(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCPublicMessageEvent,
            text=event.arguments[0],
            nick=event.source.nick,
            channel=event.target,
            mentions_me=self._mentions_me(connection, event.arguments[0]),
        )

    def on_privmsg(self, connection: ServerConnection, event: _IRCEvent):
        self._post_event(
            connection,
            IRCPrivateMessageEvent,
            text=event.arguments[0],
            nick=event.source.nick,
            channel=event.target,
            mentions_me=self._mentions_me(connection, event.arguments[0]),
        )

    def on_dccchat(self, connection: DCCConnection, event: _IRCEvent):
        if len(event.arguments) != 2:
            return

        args = event.arguments[1].split()
        if len(args) == 4:
            try:
                address = ip_numstr_to_quad(args[2])
                port = int(args[3])
            except ValueError:
                return

            nick = event.source.nick
            if not self._is_dcc_connection_request_allowed(address=address, nick=nick):
                self.logger.info(
                    f'Refused DCC connection from address={address} nick={nick}'
                )
                connection.disconnect('Unauthorized peer')
                return

            self._post_event(
                connection, IRCDCCRequestEvent, address=address, port=port, nick=nick
            )

            self.dcc('chat').connect(address, port)

    def _read(self, sock: socket.socket, bufsize: int = 4096) -> bytes:
        if self.dcc_file_transfer_timeout:
            sock.setblocking(False)
            has_data = select.select([sock], [], [], self.dcc_file_transfer_timeout)
            if has_data[0]:
                return sock.recv(bufsize)
            raise TimeoutError(f'Time out ({self.dcc_file_transfer_timeout}s)')

        return sock.recv(bufsize)

    def _dcc_connect_processor(self):
        while True:
            proc_type, addr, port = self._dcc_proc_completion_queue.get()
            conn_map = getattr(self, f'_dcc_{proc_type}_procs', {})
            conn_map.pop((addr, port), None)

    def _process_dcc_recv(
        self,
        connection: DCCConnection,
        filename: str,
        address: str,
        port: int,
        size: int,
    ):
        pathlib.Path(self.dcc_downloads_dir).mkdir(parents=True, exist_ok=True)
        filename = os.path.abspath(os.path.join(self.dcc_downloads_dir, filename))
        assert filename.startswith(
            self.dcc_downloads_dir
        ), 'Attempt to save a file outside the downloads directory'

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock, open(
            filename, 'wb'
        ) as f:
            try:
                sock.connect((address, port))
                processed_bytes = 0

                while processed_bytes < size:
                    buf = self._read(sock)
                    processed_bytes += len(buf)
                    f.write(buf)
                    sock.send(f'{socket.htonl(processed_bytes)}'.encode())

                self._post_event(
                    connection,
                    IRCDCCFileRecvCompletedEvent,
                    address=address,
                    port=port,
                    file=filename,
                    size=size,
                )
            except Exception as e:
                self.logger.error(f'DCC transfer error from {address}:{port}: {e}')
                self.logger.exception(e)
                self._post_event(
                    connection,
                    IRCDCCFileRecvCancelledEvent,
                    address=address,
                    port=port,
                    file=filename,
                    error=str(e),
                )
            finally:
                connection.disconnect()
                self._dcc_proc_completion_queue.put(('recv', address, port))

    def _process_dcc_send(self, connection: DCCConnection, filename: str):
        try:
            dcc_sock, addr = connection.socket.accept()
            self._set_accept_timeout(dcc_sock)
            self.logger.info(f'Accepted DCC connection from {addr}')

            # Create a new server socket for the file transfer
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as srv_sock:
                srv_sock.bind((connection.localaddress, 0))
                address, transfer_port = srv_sock.getsockname()
                dcc_sock.send(
                    (
                        f'CTCP_MESSAGE \x01'
                        f'DCC SEND {os.path.basename(filename)} '
                        f'{ip_quad_to_numstr(address)} {transfer_port} '
                        f'{os.path.getsize(filename)}\x01\r\n'
                    ).encode()
                )

                # Wait for a connection on the file transfer port
                self._set_accept_timeout(srv_sock)
                srv_sock.listen(1)
                client_sock, addr = srv_sock.accept()
                self.logger.info(f'DCC file transfer for {filename} accepted on {addr}')

                try:
                    with open(filename, 'rb') as f:
                        buf = f.read(4096)
                        while buf:
                            client_sock.send(buf)
                            recv_bytes = struct.unpack('>i', self._read(client_sock))[0]
                            buf = buf[recv_bytes:]

                            if not buf:
                                buf = f.read(4096)

                    self._post_event(
                        connection,
                        IRCDCCFileSendCompletedEvent,
                        file=filename,
                        address=address,
                        port=transfer_port,
                    )
                finally:
                    client_sock.close()
                    dcc_sock.close()
        except Exception as e:
            self.logger.error(
                f'DCC transfer error to {connection.peeraddress}:{connection.peerport}: {e}'
            )
            self.logger.exception(e)
            self._post_event(
                connection,
                IRCDCCFileSendCancelledEvent,
                address=connection.peeraddress,
                port=connection.peerport,
                file=filename,
                error=str(e),
            )
        finally:
            connection.disconnect()
            self._dcc_proc_completion_queue.put(
                ('send', connection.localaddress, connection.localport)
            )

    def _set_accept_timeout(self, sock: socket.socket):
        if self.dcc_accept_timeout:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.settimeout(self.dcc_accept_timeout)

    def _stop_dcc_connection(self, address: str, port: int):
        self.logger.info(f'Stopping DCC connection to {address}:{port}')
        proc = self._dcc_recv_procs.pop((address, port), None)
        if proc and proc.is_alive():
            proc.terminate()

    def _accept_dcc_file_request(
        self, connection: DCCConnection, ctcp_msg: str, nick: str
    ):
        ctcp_msg = [token.strip() for token in ctcp_msg.split(' ') if token]

        filename = ' '.join(ctcp_msg[2:-3])
        address, port, size = ctcp_msg[-3:]

        if not (filename and address and port):
            return

        address = ip_numstr_to_quad(address)
        port = int(port)
        size = int(size)

        if port == 0:
            self.logger.warning('Passive CTCP file transfer is currently not supported')
            return

        self._post_event(
            connection,
            IRCDCCFileRequestEvent,
            address=address,
            port=port,
            file=filename,
            size=size,
            nick=nick,
        )

        # Accept the file request
        # (if we're here then the peer is whitelisted/not blacklisted)
        self._dcc_recv_procs[(address, port)] = multiprocessing.Process(
            target=self._process_dcc_recv,
            kwargs={
                'connection': connection,
                'address': address,
                'port': port,
                'filename': filename,
                'size': size,
            },
        )

        self._dcc_recv_procs[(address, port)].start()

    def _handle_ctcp_message(self, connection: DCCConnection, event: _IRCEvent):
        body = event.arguments[0]
        sep_pos = [i for i, ch in enumerate(body) if ch == ord(b'\x01')]
        msg = []

        try:
            msg = [
                body[(0 if i == 0 else sep_pos[i - 1] + 1) : pos].decode().strip()
                for i, pos in enumerate(sep_pos)
            ][1:]
        except (ValueError, TypeError):
            pass

        if not msg:
            return

        # Check if this is a DCC file send request
        if msg[0].startswith('DCC SEND '):
            self._accept_dcc_file_request(connection, msg[0], event.source)
        else:
            self._post_event(
                connection, IRCCTCPMessageEvent, address=event.source, message=msg
            )

    def on_dccmsg(self, connection: DCCConnection, event: _IRCEvent):
        ctcp_header = b'CTCP_MESSAGE'
        if event.arguments[0][: len(ctcp_header)] == ctcp_header:
            return self._handle_ctcp_message(connection, event)

        self._post_event(
            connection,
            IRCDCCMessageEvent,
            address=event.source,
            body=event.arguments[0],
        )

    def on_disconnect(self, connection: Connection, event: _IRCEvent):
        self._post_event(connection, IRCDisconnectEvent)
        # Cache channels for reconnect logic
        channels = {ch: None for ch in self.channels.keys()}
        super()._on_disconnect(connection, event)
        self.channels.update(channels)

    def on_whoisuser(self, _, event: _IRCEvent):
        self._process_event_on_queue(event)

    def on_currenttopic(self, _, event: _IRCEvent):
        self._process_event_on_queue(event)

    def _process_event_on_queue(self, event: _IRCEvent):
        q = self._pending_requests.get(event.type)
        if q:
            q.put(event)

    def _on_generic_event(self, _, event: _IRCEvent):
        self.logger.debug(
            f'Received raw unhandled IRC event on {self.server}: {event.__dict__}'
        )

    def _is_dcc_connection_request_allowed(self, address: str, nick: str) -> bool:
        if self.dcc_ip_whitelist and address not in self.dcc_ip_whitelist:
            return False
        if self.dcc_ip_blacklist and address in self.dcc_ip_blacklist:
            return False
        if self.dcc_nick_whitelist and nick not in self.dcc_nick_whitelist:
            return False
        if self.dcc_nick_blacklist and nick in self.dcc_nick_blacklist:
            return False

        if (
            self.dcc_max_connections
            and len(self._dcc_recv_procs) >= self.dcc_max_connections
        ):
            self.logger.warning(
                f'Refused new DCC connection: maximum number of concurrent '
                f'connections ({self.dcc_max_connections}) reached'
            )

            return False

        return True

    def dcc_file_transfer(
        self, file: str, nick: str, bind_address: Optional[str] = None
    ):
        conn: DCCConnection = self.dcc('chat')
        bind_address = (bind_address, 0) if bind_address else None
        conn.listen(bind_address)
        conn.passive = False
        self.connection.privmsg(
            nick,
            f'\x01DCC CHAT chat {ip_quad_to_numstr(conn.localaddress)} {conn.localport}\x01',
        )

        send_proc = self._dcc_send_procs[
            (conn.localaddress, conn.localport)
        ] = multiprocessing.Process(
            target=self._process_dcc_send,
            kwargs={
                'connection': conn,
                'filename': file,
            },
        )

        send_proc.start()
        send_proc.join()

    @contextlib.contextmanager
    def event_queue(
        self, event_type: str
    ) -> Generator[multiprocessing.Queue, None, None]:
        q = self._pending_requests[event_type] = multiprocessing.Queue()
        try:
            yield q
        finally:
            q.close()
            self._pending_requests.pop(event_type, None)

    def start(self):
        self._dcc_processor = multiprocessing.Process(
            target=self._dcc_connect_processor
        )
        self._dcc_processor.start()

        for event_code in irc_event_codes.keys():
            handler = getattr(self, f'on_{event_code}', None)
            if not handler:
                self.reactor.add_global_handler(
                    event_code, lambda conn, evt: self._on_generic_event(conn, evt)
                )

        super().start()

    def stop(self, msg: Optional[str] = None):
        msg = msg or self.stop_message
        for addr, port in list(self._dcc_recv_procs.keys()) + list(
            self._dcc_send_procs.keys()
        ):
            self._stop_dcc_connection(addr, port)

        if self._dcc_processor and self._dcc_processor.is_alive():
            self._dcc_processor.terminate()
        self._dcc_processor = None

        super().disconnect(msg=msg)