import logging
import uuid

import zmq
from zmq.eventloop.zmqstream import ZMQStream


class SessionManager(object):

    def __init__(self, kernel_manager, kernel_id, context):
        self.context = context
        self.kernel_manager = kernel_manager
        self.kernel_id = kernel_id
        self._sessions = {}

    def __del__(self):
        self.stop_all()

    @property
    def session_ids(self):
        return self._session.keys()

    def __len__(self):
        return len(self.session_ids)

    def __contains__(self, session_id):
        if session_id in self.session_ids:
            return True
        else:
            return False

    def start_session(self):
        session_id = str(uuid.uuid4())
        ports = self.kernel_manager.get_kernel_ports(self.kernel_id)
        iopub_stream = self.create_connected_stream(ports['iopub_port'], zmq.SUB)
        iopub_stream.socket.setsockopt(zmq.SUBSCRIBE, b'')
        shell_stream = self.create_connected_stream(ports['shell_port'], zmq.XREQ)
        self._sessions[session_id] = dict(
            iopub_stream = iopub_stream,
            shell_stream = shell_stream
        )
        return session_id

    def stop_session(self, session_id):
        session_dict = self._sessions.get(session_id)
        if session_dict is not None:
            for name, stream in session_dict.items():
                stream.close()
            del self._sessions[session_id]

    def stop_all(self):
        for session_id in self._sessions.keys():
            self.stop_session(session_id)

    def create_connected_stream(self, port, socket_type):
        sock = self.context.socket(socket_type)
        addr = "tcp://%s:%i" % (self.kernel_manager.ip, port)
        logging.info("Connecting to: %s, %r" % (addr, socket_type))
        sock.connect(addr)
        return ZMQStream(sock)

    def get_stream(self, session_id, stream_name):
        session_dict = self._sessions.get(session_id)
        if session_dict is not None:
            return session_dict[stream_name]
        else:
            raise KeyError("Session with id not found: %s" % session_id)

    def get_iopub_stream(self, session_id):
        return self.get_stream(session_id, 'iopub_stream')

    def get_shell_stream(self, session_id):
        return self.get_stream(session_id, 'shell_stream')