from collections import defaultdict
import logging
from threading import Thread

from fuocore.daemon.thread_tcp_server import TcpServer


logger = logging.getLogger(__name__)


class DeadSubscriber(Exception):
    pass


class Subscriber(object):
    def __init__(self, addr, conn):
        self._addr = addr
        self._conn = conn

    def __eq__(self, obj):
        return self._addr == obj._addr

    def __hash__(self):
        return id(self._addr)


def sendto_subscriber(subscriber, msg):
    try:
        subscriber._conn.send(bytes(msg, 'utf-8'))
    except BrokenPipeError:
        subscriber._conn.close()
        del subscriber
        raise DeadSubscriber


class Gateway(object):
    def __init__(self):
        self.topics = set()
        self._relations = defaultdict(set)  # {'topic': subscriber_set}

    def add_topic(self, topic):
        self.topics.add(topic)

    def remove_topic(self, topic):
        if topic in self.topics:
            self.topics.remove(topic)

    def link(self, topic, subscriber):
        self._relations[topic].add(subscriber)

    def unlink(self, topic, subscriber):
        if topic in self.topics and subscriber in self._relations[topic]:
            self._relations[topic].remove(subscriber)

    def remove_subscriber(self, subscriber):
        for topic in self.topics:
            if subscriber in self._relations[topic]:
                self._relations[topic].remove(subscriber)

    def publish(self, msg, topic):
        # NOTE: use queue? maybe.
        subscribers = self._relations[topic]
        for subscriber in subscribers:
            try:
                sendto_subscriber(subscriber, msg)
            except DeadSubscriber:
                # NOTE: need lock?
                self._relations[topic].remove(subscriber)
                break


def handle(conn, addr, gateway, *args, **kwargs):
    """
    NOTE: use tcp instead of udp because some operations need ack
    """
    conn.sendall(b'OK feeluown live lyric\n')
    while True:
        try:
            s = conn.recv(1024).decode('utf-8').strip()
            if not s:
                conn.close()
                break
        except ConnectionResetError:
            logger.debug('Client close the connection.')
            break

        parts = s.split(' ')
        if len(parts) != 2:
            conn.send(b"Invalid command\n")
            continue
        cmd, topic = parts
        if cmd.lower() != 'sub':
            conn.send(bytes("Unknown command '{}'\n".format(cmd.lower()), 'utf-8'))
            continue
        if topic not in gateway.topics:
            conn.send(bytes("Unknown topic '{}'\n".format(topic), 'utf-8'))
            continue
        subscriber = Subscriber(addr, conn)
        gateway.link(topic, subscriber)
        break


def run(host='0.0.0.0', port=23334):
    gateway = Gateway()
    server = TcpServer(host, port, handle_func=handle)
    Thread(target=server.run, args=(gateway,)).start()
    logger.info('Run pubsub server in {host}:{port}'.format(
        host=host, port=port))
    return gateway, server


if __name__ == '__main__':
    import time
    gateway, server = run()
    print('pubsub is running.')
    gateway.add_topic('topic.live_lyric')
    while True:
        time.sleep(1)
        gateway.publish('miao\n', 'topic.live_lyric')
