// DWBS ------------------------------------------------------------------ DWBS
//
// PLATFORM		: ALL
// PRODUCT		: BIT_DEMON
// VISIBILITY	: PUBLIC
//
// DWBS ------------------------------------------------------------------ DWBS

#include <bdNet/bdStream/bdStreamStore.h>

#include <bdCore/bdUtilities/bdSort.h>

#include <bdConnection/bdMessage/bdMessage.h>
#include <bdConnection/bdMessage/bdReceivedMessage.h>
#include <bdConnection/bdMessage/bdMessageTypes.h>

#include <bdNet/bdDispatch/bdDispatcher.h>

#include <bdNet/bdStream/bdStreamListener.h>
#include <bdNet/bdStream/bdInStream.h>
#include <bdNet/bdStream/bdOutStream.h>

#define BD_LOG_LEVEL "bdNet/streams"

#define BD_FIRE_LISTENER(l)								\
do														\
{														\
	for(bdUInt i = 0; i < m_listeners.getSize(); i++)	\
	{													\
		m_listeners[i]->l;								\
	}													\
}														\
while(g_assertFalse)

bdStreamStore::bdStreamStore(bdDispatcher *dispatcher)
: m_okToSend(true),
  m_dispatcher(dispatcher)
{
	m_dispatcher->registerInterceptor(this);
}

bdStreamStore::~bdStreamStore()
{
	m_dispatcher->unregisterInterceptor(this);
}

bdBool bdStreamStore::registerOutStream(const bdOutStreamRef outStream)
{
	// send message to tell the other side to open
	const bdBool sent = sendCreateInStream(outStream);

	if(sent)
	{
		// add the out going stream to our list
		m_outStreams.pushBack(outStream);
		
		sortOutStreams();
	}

	return sent;
}

bdBool bdStreamStore::closeInStream(const bdInStreamRef stream)
{
	return sendCloseOutStream(stream->getStreamId().getStreamNum(), stream->getConnection());
}

void bdStreamStore::registerListener(bdStreamListener *const listener)
{
#ifdef BD_DEBUG
	bdUInt32 index;
	if(m_listeners.findFirst(listener, index))
	{
		BD_WARN(BD_LOG_LEVEL, "bdStreamStore::registerListener, same listener registered multiple times.");
	}
#endif // BD_DEBUG

	m_listeners.pushBack(listener);
}

void bdStreamStore::unregisterListener(bdStreamListener *const listener)
{
	m_listeners.removeAllKeepOrder(listener);
}

void bdStreamStore::doWork()
{
	if(m_okToSend)
	{		
		for(bdUInt i = 0; i < m_outStreams.getSize(); i++)
		{
			const bdOutStreamRef outStream = m_outStreams[i];

			// TODO: is this calculation correct?
			const bdConnectionRef connection = outStream->getConnection();
			const bdUInt transmissionRate = connection->getTransmissionRate();			
			const bdUInt bytesSent = connection->getStats().getBytesSentPerSecond();

			if(outStream->isFinished())
			{
				sendFinishedInStream(outStream);
			}
			else if(outStream->isClosed())
			{
				sendCloseInStream(outStream);
			}
			else if ((bytesSent <= transmissionRate) && outStream->hasData())
			{
				const bdMessageRef message = new bdMessage(BD_STREAM_DATA);
				const bdBitBufferRef buffer = message->getPayload();
				const bdBitBufferRef unencBuffer = message->getUnencryptedPayload();

				// write the stream number
				buffer->writeByte8(outStream->getNum());

				// TODO: is this calculation correct?
				const bdUInt maxNumBytes = transmissionRate - bytesSent;
				if(outStream->getData(buffer, unencBuffer, maxNumBytes))
				{
					connection->send(message, outStream->getReliable());
				}
			}
		}
	}
}

bdBool bdStreamStore::accept(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdMessageType type = message->getType();
	bdBool accepted = true;

	switch(type)
	{
		case BD_STREAM_CREATE:
		{
			handleStreamCreate(receivedMsg);
		}
		break;		
		case BD_STREAM_OUT_STREAM_CLOSE:
		{
			handleOutStreamClose(receivedMsg);
		}
		break;
		case BD_STREAM_IN_STREAM_FINISHED:
		{
			handleInStreamFinished(receivedMsg);
		}
		case BD_STREAM_IN_STREAM_CLOSE:
		{
			handleInStreamClose(receivedMsg);
		}
		break;
		case BD_STREAM_DATA:
		{
			handleStreamData(receivedMsg);
		}
		break;
		default:
		{
			accepted = false;
		}
	}

	return accepted;
}

void bdStreamStore::handleStreamCreate(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdBitBufferRef buffer = message->getPayload();

	// read the stream number
	bdStreamNum streamNum;
	bdBool ok = buffer->readUByte8(streamNum);

	// read the stream type
	bdStreamType streamType = 0;
	ok = ok && buffer->readUByte8(streamType);

	if(ok)
	{
		const bdConnectionRef connection = receivedMsg.getConnection();
		const bdCommonAddrRef addr = connection->getAddress();
		const bdStreamId streamId(addr, streamNum);

		bdInStreamRef instream;
		for(bdUInt i = 0; instream.isNull() && (i < m_listeners.getSize()); i++)
		{
			instream = m_listeners[i]->onNewInStream(streamType, streamId, connection);			
		}

		if(instream.isNull())
		{
			BD_WARN(BD_LOG_LEVEL, "No stream listeners registered!");

			// none of the listeners want to handle this stream so tell the 
			// other end to stop sending
			sendCloseOutStream(streamNum, connection);
		}
		else
		{
			maintainListener(connection, true);

			if(!m_inStreams.put(streamId, instream))
			{
				BD_WARN(BD_LOG_LEVEL, "Unable to add new stream to list. Closing.");

				sendCloseOutStream(streamNum, connection);
			}
		}
	}
}

void bdStreamStore::handleOutStreamClose(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdBitBufferRef buffer = message->getPayload();

	// read the stream number
	bdStreamNum streamNum;
	bdBool ok = buffer->readUByte8(streamNum);

	if(ok)
	{
		// find the out stream
		bdBool removed = false;
		const bdUInt numStreams = m_outStreams.getSize();
		for(bdUInt i = 0; !removed && (i < numStreams); i++)
		{
			const bdOutStreamRef stream = m_outStreams[i];
			if(stream->getNum() == streamNum)
			{
				if(stream->getConnection() == receivedMsg.getConnection())
				{
					maintainListener(stream->getConnection(), false);
					BD_FIRE_LISTENER(onOutStreamClose(stream));
					m_outStreams.removeSection(i, i+1);
					removed = true;
				}
				else
				{
					BD_WARN(BD_LOG_LEVEL, "Close came in from wrong connection.");
					break;
				}
			}
		}

		if(!removed)
		{
			BD_WARN(BD_LOG_LEVEL, "bdStreamStore::handleOutStreamClose, could not find stream to close.");
		}
	}
}


void bdStreamStore::handleInStreamFinished(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdBitBufferRef buffer = message->getPayload();

	// read the stream number
	bdStreamNum streamNum;
	bdBool ok = buffer->readUByte8(streamNum);

	if(ok)
	{		
		const bdCommonAddrRef addr = receivedMsg.getConnection()->getAddress();
		const bdStreamId streamId(addr, streamNum);

		// find the instream
		bdInStreamRef inStream;
		if(m_inStreams.get(streamId, inStream))
		{
			if(inStream->getConnection() == receivedMsg.getConnection())
			{
				maintainListener(inStream->getConnection(), false);
				BD_FIRE_LISTENER(onInStreamFinished(inStream));
				// remove the in stream
				m_inStreams.remove(streamId);
			}
			else
			{
				BD_WARN(BD_LOG_LEVEL, "Finished from wrong connection!");
			}
		}
	}
}


void bdStreamStore::handleInStreamClose(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdBitBufferRef buffer = message->getPayload();

	// read the stream number
	bdStreamNum streamNum;
	bdBool ok = buffer->readUByte8(streamNum);

	if(ok)
	{		
		const bdCommonAddrRef addr = receivedMsg.getConnection()->getAddress();
		const bdStreamId streamId(addr, streamNum);

		// find the instream
		bdInStreamRef inStream;
		if(m_inStreams.get(streamId, inStream))
		{
			if(inStream->getConnection() == receivedMsg.getConnection())
			{
				maintainListener(inStream->getConnection(), false);
				BD_FIRE_LISTENER(onInStreamClose(inStream));
				// remove the in stream
				m_inStreams.remove(streamId);
			}
			else
			{
				BD_WARN(BD_LOG_LEVEL, "Finished from wrong connection!");
			}
		}
	}
}

void bdStreamStore::handleStreamData(const bdReceivedMessage& receivedMsg)
{
	const bdMessageRef message = receivedMsg.getMessage();
	const bdBitBufferRef buffer = message->getPayload();
	const bdBitBufferRef unencBuffer = message->getUnencryptedPayload();
	
	// read the stream number
	bdStreamNum streamNum;
	bdBool ok = buffer->readUByte8(streamNum);

	if(ok)
	{		
		const bdCommonAddrRef addr = receivedMsg.getConnection()->getAddress();
		const bdStreamId streamId(addr, streamNum);		

		// find the instream
		bdInStreamRef inStream;
		if(m_inStreams.get(streamId, inStream))
		{
			if(inStream->getConnection() == receivedMsg.getConnection())
			{
				inStream->receiveData(buffer, unencBuffer);
			}
			else
			{
				BD_WARN(BD_LOG_LEVEL, "Data from wrong connection!");
			}
		}
	}
}

bdBool bdStreamStore::sendCreateInStream(const bdOutStreamRef outStream)
{
	const bdMessageRef createMessage = new bdMessage(BD_STREAM_CREATE);
	const bdBitBufferRef buffer = createMessage->getPayload();

	buffer->writeUByte8(outStream->getNum());
	buffer->writeUByte8(outStream->getType());

	const bdConnectionRef connection = outStream->getConnection();
	const bdBool sent = connection->send(createMessage, true);

	// register as a connection listener so we can close the
	// stream if we disconnect
	maintainListener(connection, true);

	return sent;
}

bdBool bdStreamStore::sendFinishedInStream(const bdOutStreamRef outStream)
{
	const bdMessageRef closeMessage = new bdMessage(BD_STREAM_IN_STREAM_FINISHED);
	const bdBitBufferRef buffer = closeMessage->getPayload();

	buffer->writeUByte8(outStream->getNum());	

	const bdConnectionRef connection = outStream->getConnection();
	const bdBool sent = connection->send(closeMessage, true);

	maintainListener(connection, false);

	return sent;
}

bdBool bdStreamStore::sendCloseInStream(const bdOutStreamRef outStream)
{
	const bdMessageRef closeMessage = new bdMessage(BD_STREAM_IN_STREAM_CLOSE);
	const bdBitBufferRef buffer = closeMessage->getPayload();

	buffer->writeUByte8(outStream->getNum());	
	
	const bdConnectionRef connection = outStream->getConnection();
	const bdBool sent = connection->send(closeMessage, true);

	maintainListener(connection, false);

	return sent;
}

bdBool bdStreamStore::sendCloseOutStream(const bdStreamType streamNum,
										 const bdConnectionRef connection)
{
	const bdMessageRef closeMessage = new bdMessage(BD_STREAM_OUT_STREAM_CLOSE);
	const bdBitBufferRef buffer = closeMessage->getPayload();

	buffer->writeUByte8(streamNum);	

	const bdBool sent = connection->send(closeMessage, true);

	maintainListener(connection, false);

	return sent;
}

void bdStreamStore::sortOutStreams()
{
	bdCombSort<bdOutStreamRef>(m_outStreams.begin(), m_outStreams.getSize(), bdOutStreamRefLessThan());
}

void bdStreamStore::maintainListener(const bdConnectionRef connection,
									 const bdBool registerListener)
{
	// have to be careful that we don't register as a listener for the 
	// same connection more than once, or unregister before all streams
	// are closed
	bdBool listening = false;

	// check for out streams to this connection
	const bdUInt numOutStreams = m_outStreams.getSize();
	for(bdUInt i=0; !listening && (i < numOutStreams); i++)
	{
		if(m_outStreams[i]->getConnection() == connection)
		{
			listening = true;
		}
	}

	// check for in streams from this connection
	bdHashMap<bdStreamId, bdInStreamRef, bdStreamIdHash>::Iterator iter = m_inStreams.getIterator();

	while(iter && !listening)
	{
		const bdInStreamRef stream = m_inStreams.getValue(iter);
		m_inStreams.next(iter);

		if(stream->getConnection() == connection)
		{
			listening = true;
		}
	}

	if(!listening)
	{
		if(registerListener)
		{
			connection->registerListener(this);
		}
		else
		{
			connection->unregisterListener(this);
		}
	}
}

bdBool bdStreamStore::bdOutStreamRefLessThan::operator() (const bdOutStreamRef& x, 
														  const bdOutStreamRef& y) const
{
	return x->getPriority() < y->getPriority();
}


void bdStreamStore::onDisconnect(bdConnectionRef connection)
{
	// check for out streams going to this connection
	for(bdUInt i = 0; i < m_outStreams.getSize(); i++)
	{
		const bdOutStreamRef stream = m_outStreams[i];
		if(stream->getConnection() == connection)
		{
			BD_FIRE_LISTENER(onOutStreamClose(stream));

			// remove the stream
			m_outStreams.removeSection(i, i+1);

			// need to recheck the same slot
			i--;
		}
	}

	// check for in streams from this connection
	bdHashMap<bdStreamId, bdInStreamRef, bdStreamIdHash>::Iterator iter = m_inStreams.getIterator();

	while(iter)
	{
		const bdInStreamRef stream = m_inStreams.getValue(iter);
		
		if(stream->getConnection() == connection)
		{
			BD_FIRE_LISTENER(onInStreamClose(stream));

			m_inStreams.remove(iter);
		}
		else
		{
			m_inStreams.next(iter);
		}
	}
}

void bdStreamStore::onSendHighWaterMark(bdUInt32 /*bytesSent*/)
{
	m_okToSend = false;
}

void bdStreamStore::onSendLowWaterMark(bdUInt32 /*bytesSent*/)
{
	m_okToSend = true;
}
