###############################################################################
##
##  Copyright (C) 2011-2014 Tavendo GmbH
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##      http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.
##
###############################################################################

from __future__ import absolute_import

__all__ = (
   'WebSocketAdapterProtocol',
   'WebSocketServerProtocol',
   'WebSocketClientProtocol',
   'WebSocketAdapterFactory',
   'WebSocketServerFactory',
   'WebSocketClientFactory',

   'WrappingWebSocketAdapter',
   'WrappingWebSocketServerProtocol',
   'WrappingWebSocketClientProtocol',
   'WrappingWebSocketServerFactory',
   'WrappingWebSocketClientFactory',

   'listenWS',
   'connectWS',

   'WampWebSocketServerProtocol',
   'WampWebSocketServerFactory',
   'WampWebSocketClientProtocol',
   'WampWebSocketClientFactory',
)


from base64 import b64encode, b64decode

from zope.interface import implementer

import twisted.internet.protocol
from twisted.internet.defer import maybeDeferred
from twisted.python import log
from twisted.internet.interfaces import ITransport

from autobahn.wamp import websocket
from autobahn.websocket import protocol
from autobahn.websocket import http

from autobahn.websocket.compress import PerMessageDeflateOffer, \
                                        PerMessageDeflateOfferAccept, \
                                        PerMessageDeflateResponse, \
                                        PerMessageDeflateResponseAccept


class WebSocketAdapterProtocol(twisted.internet.protocol.Protocol):
   """
   Adapter class for Twisted WebSocket client and server protocols.
   """

   def connectionMade(self):
      ## the peer we are connected to
      try:
         peer = self.transport.getPeer()
      except:
         ## ProcessProtocols lack getPeer()
         self.peer = "?"
      else:
         try:
            self.peer = "%s:%d" % (peer.host, peer.port)
         except:
            ## eg Unix Domain sockets don't have host/port
            self.peer = str(peer)

      self._connectionMade()

      ## Set "Nagle"
      try:
         self.transport.setTcpNoDelay(self.tcpNoDelay)
      except:
         ## eg Unix Domain sockets throw Errno 22 on this
         pass


   def connectionLost(self, reason):
      self._connectionLost(reason)


   def dataReceived(self, data):
      self._dataReceived(data)


   def _closeConnection(self, abort = False):
      if abort and hasattr(self.transport, 'abortConnection'):
         ## ProcessProtocol lacks abortConnection()
         self.transport.abortConnection()
      else:
         self.transport.loseConnection()


   def _onOpen(self):
      self.onOpen()

 
   def _onMessageBegin(self, isBinary):
      self.onMessageBegin(isBinary)

 
   def _onMessageFrameBegin(self, length):
      self.onMessageFrameBegin(length)

 
   def _onMessageFrameData(self, payload):
      self.onMessageFrameData(payload)

 
   def _onMessageFrameEnd(self):
      self.onMessageFrameEnd()


   def _onMessageFrame(self, payload):
      self.onMessageFrame(payload)


   def _onMessageEnd(self):
      self.onMessageEnd()


   def _onMessage(self, payload, isBinary):
      self.onMessage(payload, isBinary)


   def _onPing(self, payload):
      self.onPing(payload)


   def _onPong(self, payload):
      self.onPong(payload)


   def _onClose(self, wasClean, code, reason):
      self.onClose(wasClean, code, reason)


   def registerProducer(self, producer, streaming):
      """
      Register a Twisted producer with this protocol.

      Modes: Hybi, Hixie

      :param producer: A Twisted push or pull producer.
      :type producer: object
      :param streaming: Producer type.
      :type streaming: bool
      """
      self.transport.registerProducer(producer, streaming)



class WebSocketServerProtocol(WebSocketAdapterProtocol, protocol.WebSocketServerProtocol):
   """
   Base class for Twisted-based WebSocket server protocols.
   """

   def _onConnect(self, request):
      ## onConnect() will return the selected subprotocol or None
      ## or a pair (protocol, headers) or raise an HttpException
      ##
      res = maybeDeferred(self.onConnect, request)

      res.addCallback(self.succeedHandshake)

      def forwardError(failure):
         if failure.check(http.HttpException):
            return self.failHandshake(failure.value.reason, failure.value.code)
         else:
            if self.debug:
               self.factory._log("Unexpected exception in onConnect ['%s']" % failure.value)
            return self.failHandshake(http.INTERNAL_SERVER_ERROR[1], http.INTERNAL_SERVER_ERROR[0])

      res.addErrback(forwardError)



class WebSocketClientProtocol(WebSocketAdapterProtocol, protocol.WebSocketClientProtocol):
   """
   Base class for Twisted-based WebSocket client protocols.
   """

   def _onConnect(self, response):
      self.onConnect(response)



class WebSocketAdapterFactory:
   """
   Adapter class for Twisted-based WebSocket client and server factories.
   """

   def _log(self, msg):
      log.msg(msg)


   def _callLater(self, delay, fun):
      return self.reactor.callLater(delay, fun)



class WebSocketServerFactory(WebSocketAdapterFactory, protocol.WebSocketServerFactory, twisted.internet.protocol.ServerFactory):
   """
   Base class for Twisted-based WebSocket server factories.

   .. seealso:: `twisted.internet.protocol.ServerFactory <http://twistedmatrix.com/documents/current/api/twisted.internet.protocol.ServerFactory.html>`_
   """

   def __init__(self, *args, **kwargs):
      """
      In addition to all arguments to the constructor of
      :class:`autobahn.websocket.protocol.WebSocketServerFactory`,
      you can supply a `reactor` keyword argument to specify the
      Twisted reactor to be used.
      """
      ## lazy import to avoid reactor install upon module import
      if 'reactor' in kwargs:
         if kwargs['reactor']:
            self.reactor = kwargs['reactor']
         else:
            from twisted.internet import reactor
            self.reactor = reactor
         del kwargs['reactor']
      else:
         from twisted.internet import reactor
         self.reactor = reactor

      protocol.WebSocketServerFactory.__init__(self, *args, **kwargs)



class WebSocketClientFactory(WebSocketAdapterFactory, protocol.WebSocketClientFactory, twisted.internet.protocol.ClientFactory):
   """
   Base class for Twisted-based WebSocket client factories.

   .. seealso:: `twisted.internet.protocol.ClientFactory <http://twistedmatrix.com/documents/current/api/twisted.internet.protocol.ClientFactory.html>`_
   """

   def __init__(self, *args, **kwargs):
      """
      In addition to all arguments to the constructor of
      :class:`autobahn.websocket.protocol.WebSocketClientFactory`,
      you can supply a `reactor` keyword argument to specify the
      Twisted reactor to be used.
      """
      ## lazy import to avoid reactor install upon module import
      if 'reactor' in kwargs:
         if kwargs['reactor']:
            self.reactor = kwargs['reactor']
         else:
            from twisted.internet import reactor
            self.reactor = reactor
         del kwargs['reactor']
      else:
         from twisted.internet import reactor
         self.reactor = reactor

      protocol.WebSocketClientFactory.__init__(self, *args, **kwargs)



@implementer(ITransport)
class WrappingWebSocketAdapter:
   """
   An adapter for stream-based transport over WebSocket.

   This follows `websockify <https://github.com/kanaka/websockify>`_
   and should be compatible with that.

   It uses WebSocket subprotocol negotiation and supports the
   following WebSocket subprotocols:

     - ``binary`` (or a compatible subprotocol)
     - ``base64``

   Octets are either transmitted as the payload of WebSocket binary
   messages when using the ``binary`` subprotocol (or an alternative
   binary compatible subprotocol), or encoded with Base64 and then
   transmitted as the payload of WebSocket text messages when using
   the ``base64`` subprotocol.
   """

   def onConnect(self, requestOrResponse):

      ## Negotiate either the 'binary' or the 'base64' WebSocket subprotocol
      ##
      if isinstance(requestOrResponse, protocol.ConnectionRequest):
         request = requestOrResponse
         for p in request.protocols:
            if p in self.factory._subprotocols:
               self._binaryMode = (p != 'base64')
               return p
         raise http.HttpException(http.NOT_ACCEPTABLE[0], "this server only speaks %s WebSocket subprotocols" % self.factory._subprotocols)
      elif isinstance(requestOrResponse, protocol.ConnectionResponse):
         response = requestOrResponse
         if response.protocol not in self.factory._subprotocols:
            self.failConnection(protocol.WebSocketProtocol.CLOSE_STATUS_CODE_PROTOCOL_ERROR, "this client only speaks %s WebSocket subprotocols" % self.factory._subprotocols)
         self._binaryMode = (response.protocol != 'base64')
      else:
         ## should not arrive here
         raise Exception("logic error")


   def onOpen(self):
      self._proto.connectionMade()


   def onMessage(self, payload, isBinary):
      if isBinary != self._binaryMode:
         self.failConnection(protocol.WebSocketProtocol.CLOSE_STATUS_CODE_UNSUPPORTED_DATA, "message payload type does not match the negotiated subprotocol")
      else:
         if not isBinary:
            try:
               payload = b64decode(payload)
            except Exception as e:
               self.failConnection(protocol.WebSocketProtocol.CLOSE_STATUS_CODE_INVALID_PAYLOAD, "message payload base64 decoding error: {}".format(e))
         #print("forwarding payload: {}".format(binascii.hexlify(payload)))
         self._proto.dataReceived(payload)


   def onClose(self, wasClean, code, reason):
      self._proto.connectionLost(None)


   def write(self, data):
      #print("sending payload: {}".format(binascii.hexlify(data)))
      ## part of ITransport
      assert(type(data) == bytes)
      if self._binaryMode:
         self.sendMessage(data, isBinary = True)
      else:
         data = b64encode(data)
         self.sendMessage(data, isBinary = False)


   def writeSequence(self, data):
      ## part of ITransport
      for d in data:
         self.write(d)


   def loseConnection(self):
      ## part of ITransport
      self.sendClose()


   def getPeer(self):
      ## part of ITransport
      return self.transport.getPeer()


   def getHost(self):
      ## part of ITransport
      return self.transport.getHost()



class WrappingWebSocketServerProtocol(WrappingWebSocketAdapter, WebSocketServerProtocol):
   """
   Server protocol for stream-based transport over WebSocket.
   """



class WrappingWebSocketClientProtocol(WrappingWebSocketAdapter, WebSocketClientProtocol):
   """
   Client protocol for stream-based transport over WebSocket.