"""
twistedchat.py

A TCP chat server in the style of 'Chatroulette' using Twisted.

Part of a study in concurrency and networking in Python:
https://github.com/mjwestcott/chatroulette in which I create versions of this
server using asyncio, gevent, Tornado, and Twisted.

Some 'features':
    - on connection, clients are prompted for their name, which will prefix all
      of their sent messages;
    - the server will notify clients when they are matched with a partner or
      their partner disconnects;
    - clients whose partner disconnects will be put back in a waiting list to
      be matched again;
    - clients in the waiting list will periodically be sent a Nietzsche aphorism
      to keep them busy;
    - clients can issue the following commands:
        /help -> describes the service, including the commands below
        /quit -> close the connection
        /next -> end current chat and wait for a new partner
    - the /next command will remember the rejected partners so as not to match
      them together again.

Clients are expected to connect via telnet.
"""
from twisted.protocols.basic import LineReceiver
from twisted.internet.protocol import Factory
from twisted.internet import reactor
import random

HOST = 'localhost'
PORT = 12345

with open("nietzsche.txt", "r") as f:
    # We'll use these to keep waiting clients busy.
    aphorisms = list(filter(bool, f.read().split("\n")))

def main():
    reactor.listenTCP(PORT, ChatFactory(), interface=HOST)
    reactor.run()

class ChatProtocol(LineReceiver):
    delimiter = b'\n'

    def __init__(self):
        self.name = None
        self.partner = None
        self.rejected = set()

    def connectionMade(self):
        # Prompt user for name. Answer will be picked up in lineReceived.
        self.sendLine(
            b'Server: Welcome to TCP chat roulette! '
            b'You will be matched with a partner.')
        self.sendLine(b'Server: What is your name?')

    def connectionLost(self, exception):
        if self.name in self.factory.clients:
            del self.factory.clients[self.name]

        # Notify the client's partner, if any.
        other = self.partner
        if other is not None:
            other.partnerDisconnected()

    def lineReceived(self, data):
        if self.name is None: # First interaction with user; data is user's name.
            self.handleName(data)
        elif data.startswith(b'/'):
            self.handleCmd(data)
        else:
            self.messagePartner(data)

    def handleName(self, name):
        if name in self.factory.clients:
            self.sendLine(b'Server: Sorry, that name is taken. Please choose again.')
        else:
            self.name = name
            self.factory.clients[self.name] = self
            self.sendLine(b'Server: Hello, %b. Please wait for a partner.' % self.name)

            # Successful onboarding; match client.
            reactor.callLater(0, self.match)

    def handleCmd(self, cmd):
        if cmd.startswith(b'/help'):
            self.sendLine(b'Server: Welcome to TCP chat roulette! You will be matched with a partner.')
            self.sendLine(b'\t/help -> display this help message')
            self.sendLine(b'\t/quit -> close the connection')
            self.sendLine(b'\t/next -> end current chat and wait for a new random partner')
        elif cmd.startswith(b'/quit'):
            self.transport.loseConnection()
        elif cmd.startswith(b'/next'):
            other = self.partner
            if other is None:
                # Command issued when not enagaged in chat with a partner.
                self.sendLine(b'Server: Sorry, no partner. Please wait.')
            else:
                self.sendLine(b'Server: Chat over. Please wait for a new partner.')
                self.rejected.add(other)
                self.partner = None
                reactor.callLater(0, self.match)

                # Let down the partner gently.
                other.partnerDisconnected()
        else:
            self.sendLine(b'Server: Command not recognised.')

    def partnerDisconnected(self):
        self.partner = None
        self.sendLine(b'Server: Partner disconnected. Please wait.')
        reactor.callLater(0, self.match)

    def messagePartner(self, msg):
        """Send msg from the sender to their partner. Prefix the message with the
        sender's name."""
        assert isinstance(msg, bytes)
        partner = self.partner
        if partner is None:
            self.sendLine(b'Server: Sorry, no partner. Please wait.')
        else:
            partner.sendLine(b'%b: %b' % (self.name, msg))

    def match(self, tries=1):
        # The global clients dict and waiting set.
        waiting = self.factory.waiting
        clients = self.factory.clients

        # Find any clients who do not have a partner and add them to the
        # waiting set.
        waiting.update(c for c in clients.values() if c.partner is None)

        # Find any clients in the waiting set who have disconnected and remove
        # them from the waiting set. (If they are disconnected they will have
        # been removed from the client list.)
        waiting.intersection_update(clients.values())

        if self not in waiting:
            # We've been matched by our partner or we've disconnected.
            return

        if len(waiting) >= 2:
            # Attempt to match clients.
            A = self
            wanted = waiting - A.rejected
            partners = [B for B in wanted if A not in B.rejected and A != B]
            if partners:
                # Match succeeded.
                B = partners.pop()
                waiting.remove(A)
                waiting.remove(B)
                A.partner = B
                B.partner = A
                A.sendLine(b'Server: Partner found! Say hello.')
                B.sendLine(b'Server: Partner found! Say hello.')
                return

        # Match failed. Periodically send something interesting.
        if tries % 5 == 0:
            aphorism = random.choice(aphorisms)
            self.transport.write(
                b'Server: Thanks for waiting! Here\'s Nietzsche:\n'
                b'\n%b\n\n' % aphorism.encode("utf-8"))

        # Exponential backoff up to a maximum sleep of 20 secs.
        reactor.callLater(min(20, (tries**2)/4), self.match, tries+1)

class ChatFactory(Factory):
    protocol = ChatProtocol

    def __init__(self):
        self.clients = {}     # From names (bytes) to ChatProtocol instances.
        self.waiting = set()  # Clients without a partner i.e. waiting to chat.

if __name__ == '__main__':
    main()
