import asyncio
import logging
import ssl

#

L = logging.getLogger(__name__)

#


class Stream(object):
	'''
	This object represent a client connection.
	It is unencrypted STREAM socket.
	'''

	def __init__(self, loop, socket, outbound_queue=None):
		self.Loop = loop
		self.Socket = socket
		if outbound_queue is None:
			self.OutboundQueue = asyncio.Queue()
		else:
			self.OutboundQueue = outbound_queue

	async def recv_into(self, buf):
		return await self.Loop.sock_recv_into(self.Socket, buf)


	def send(self, data):
		self.OutboundQueue.put_nowait(data)


	async def outbound(self):
		'''
		Handle outbound direction
		'''
		while True:
			data = await self.OutboundQueue.get()
			if data is None:
				break
			await self.Loop.sock_sendall(self.Socket, data)


	async def close(self):
		self.Socket.close()


class TLSStream(object):
	'''
	This object represent a TLS client connection.
	It is encrypted SSL/TLS socket abstraction.
	'''

	def __init__(self, loop, sslcontext, socket, server_side: bool):
		self.Loop = loop
		self.Socket = socket

		self.SSLContext = sslcontext

		self.InBuffer = ssl.MemoryBIO()
		self.OutBuffer = ssl.MemoryBIO()
		self.OutboundQueue = asyncio.Queue()

		self.SSLSocket = sslcontext.wrap_bio(self.InBuffer, self.OutBuffer, server_side=server_side)


	async def handshake(self):
		while True:

			try:
				self.SSLSocket.do_handshake()
				break

			except ssl.SSLWantReadError:

				if self.OutBuffer.pending > 0:
					data = self.OutBuffer.read()
					await self.Loop.sock_sendall(self.Socket, data)

				data = await self.Loop.sock_recv(self.Socket, 4096)
				if len(data) == 0:
					# Socket has been closed
					# TODO: self.Socket.shutdown()
					return False

				self.InBuffer.write(data)
				continue

			except Exception:
				L.exception("SSL handshake failed")
				self.Socket.close()
				return False

		return True


	async def recv_into(self, buf):
		while True:

			try:
				data = self.SSLSocket.read(len(buf))
				if len(data) == 0:
					# Socket has been closed
					return 0

				buf[:len(data)] = data
				return len(data)

			except ssl.SSLWantReadError:

				if self.OutBuffer.pending > 0:
					data = self.OutBuffer.read()
					await self.Loop.sock_sendall(self.Socket, data)

				data = await self.Loop.sock_recv(self.Socket, 4096)
				if len(data) == 0:
					# Socket has been closed
					# TODO: self.Socket.shutdown()
					return 0

				self.InBuffer.write(data)
				continue


	def send(self, data):
		self.OutboundQueue.put_nowait(data)


	async def outbound(self):
		'''
		Handle outbound direction
		'''
		while True:
			data = await self.OutboundQueue.get()
			if data is None:
				break

			self.SSLSocket.write(data)

			while self.OutBuffer.pending > 0:
				data = self.OutBuffer.read()
				await self.Loop.sock_sendall(self.Socket, data)


	async def close(self):
		self.OutboundQueue.put_nowait(None)
		self.Socket.close()
