# -*- coding: utf-8 -*-

from Queue import Queue
from IPython.zmq.session import Session, Message, extract_header
from IPython.utils.traitlets import Type, HasTraits, TraitType
from IPython.zmq.kernelmanager import KernelManager, SubSocketChannel, \
XReqSocketChannel, RepSocketChannel, HBSocketChannel
MetaHasTraits = type(HasTraits)


class SubSocketChannel2p(SubSocketChannel):
    #---------------------------------------------------------------------------
    # 'SubSocketChannel' interface
    #---------------------------------------------------------------------------
    _msg = None
    queue = Queue(-1)
    def call_handlers(self, msg):  
      self.queue.put(Message(msg))
        
    def get_msg(self):
      return self.queue.get()
    
    def was_called(self):
      return not self.queue.empty()

class XReqSocketChannel2p(XReqSocketChannel):
   #---------------------------------------------------------------------------
   # 'XReqSocketChannel' interface
   #---------------------------------------------------------------------------
    _msg = None
    _called = False
    queue = Queue(-1)
    def call_handlers(self, msg):
        self.queue.put(Message(msg))
        
    def get_msg(self):
      return self.queue.get()
    
    def was_called(self):
      return not self.queue.empty()
       
class RepSocketChannel2p(RepSocketChannel):
   #---------------------------------------------------------------------------
   # 'XReqSocketChannel' interface
   #---------------------------------------------------------------------------
    _msg = None
    _called = False
    def call_handlers(self, msg):
        self._called = True
        self._msg = Message(msg)
        
    def get_msg(self):
       self._called = False
       return self._msg
    
    def was_called(self):
       return self._called
   
class HBSocketChannel2p(HBSocketChannel):
   #---------------------------------------------------------------------------
   # 'XReqSocketChannel' interface
   #---------------------------------------------------------------------------
    _msg = None
    _called = False
    def call_handlers(self, msg):
        self._called = True
        self._msg = Message(msg)
        
    def get_msg(self):
       self._called = False
       return self._msg
    
    def was_called(self):
       return self._called

class KernelManager2p(KernelManager):
      sub_channel_class = Type(SubSocketChannel2p)
      xreq_channel_class = Type(XReqSocketChannel2p)
      rep_channel_class = Type(RepSocketChannel2p)
      hb_channel_class = Type(HBSocketChannel2p)
      
      def start_kernel(self, *args, **kw):
          """ Reimplemented for proper heartbeat management.
          """
          if self._xreq_channel is not None:
              self._xreq_channel.reset_first_reply()
          super(KernelManager2p, self).start_kernel(*args, **kw)