#include "client.h"
#include "receivefile.h"
#include "sock.h"
#include "config.h"
#include "hashcache.h"
#include "tempqueue.h"
#include <map>
#include <time.h>

typedef std::pair<unsigned, unsigned short> TAddr;
static const time_t TIMEOUT_SEND = 2;

class CClientKeepAlive : public IHashProgressCallback
{
public:
	CClientKeepAlive( TAddr addr, CSock& sock ) : m_addr(addr), m_sock(sock) {}

	void OnProgress()
	{
		std::cout << "Send keep-alive" << std::endl;
		BOutputStream out;
		out << 'K' << PROTOCOL_VERSION;
		m_sock.Send( m_addr.first, m_addr.second, out.GetPtr(), out.GetSize() );
	}

private:
	TAddr m_addr;
	CSock& m_sock;
};

class CClient
{
public:
	CClient( IClientListener * pListener );

	void Run();

private:
	typedef Ptr<CReceiveFile> RFP;
	typedef std::map<TAddr, RFP> RFPMap;

	void UpdateTimeout();
	void CheckReceive();
	void CheckTimeouts();
	void CheckSend();
	void HandleDataPacket( TAddr addr, BInputStream& input );
	void HandleAnnouncePacket( TAddr addr, BInputStream& input );
	void HandleFinishPacket( TAddr addr, BInputStream& input );
	void HandleQuitPacket( TAddr addr, BInputStream& input );
	void SendRequestAnnouncement( TAddr addr );
	void SendRequestData( TAddr addr, RFP ptr, bool islandsOnly, bool throttle );
	void UpdateProgressMessage();

	RFPMap m_receiving;
	std::auto_ptr<timeval> m_pTime;
	CSock	m_sock;
	time_t m_tmSend;
	CTempQueue<TAddr> m_announceRequestThrottle;
	CTempQueue<TAddr> m_dataRequestThrottle;
	CTempQueue<TAddr> m_timeoutConnectionQueue;

	typedef void (CClient::* HandlePacket)( TAddr addr, BInputStream& input );
	typedef std::map<char, HandlePacket> TPacketHandlers;
	TPacketHandlers m_packetHandlers;
	IClientListener * m_pListener;
	bool m_bDone;
};

CClient::CClient( IClientListener * pListener ) : m_tmSend(0), m_pListener(pListener), m_bDone(false)
{
	m_sock.Bind( CLIENT_PORT );

	m_packetHandlers['D'] = &CClient::HandleDataPacket;
	m_packetHandlers['A'] = &CClient::HandleAnnouncePacket;
	m_packetHandlers['F'] = &CClient::HandleFinishPacket;
	m_packetHandlers['Q'] = &CClient::HandleQuitPacket;
}

void CClient::Run()
{
	while (!m_bDone)
	{
		UpdateTimeout();
		CheckReceive();
		CheckTimeouts();
		CheckSend();
	}
}

void CClient::UpdateTimeout()
{
	if (m_receiving.empty())
	{
		m_pTime.reset();
//		if (m_pListener)
//			m_pListener->Notification("Debug", "Set timeout infinite");
	}
	else if (!m_pTime.get())
	{
		m_pTime.reset( new timeval );
		m_pTime->tv_sec = 1;
		m_pTime->tv_usec = 0;
//		if (m_pListener)
//			m_pListener->Notification("Debug", "Set timeout 1");
	}
}

void CClient::CheckReceive()
{
	if (m_sock.Select( m_pTime.get(), true, false ))
	{
		static const size_t LENGTH = 65536;
		char buf[LENGTH];
		TAddr addr;
		size_t length = m_sock.Receive( addr.first, addr.second, buf, LENGTH );
		BInputStream input(buf, length);
		try
		{
			char packetType; input >> packetType;
			unsigned short version; input >> version;
			if (version != PROTOCOL_VERSION)
			{
				char buf[512];
				sprintf( buf, "Invalid protocol version %d (should be %d)", version, PROTOCOL_VERSION );
				throw std::runtime_error(buf);
			}
			TPacketHandlers::iterator iter = m_packetHandlers.find( packetType );
			if (iter == m_packetHandlers.end())
			{
				std::cout << "Ignored packet of type " << packetType << std::endl;
			}
			else
			{
				m_timeoutConnectionQueue.Add( 10, addr );
				HandlePacket func = iter->second;
				(this->*func)( addr, input );
			}
		}
		catch (std::exception& e)
		{
			std::cout << "Exception occured on input: " << e.what() << std::endl;
		}
	}
}

void CClient::CheckSend()
{
	if (!m_receiving.empty() && (time(NULL) - m_tmSend) > TIMEOUT_SEND)
	{
		m_tmSend = time(NULL);

		for (RFPMap::iterator iter = m_receiving.begin(); iter != m_receiving.end(); ++iter)
			SendRequestData( iter->first, iter->second, true, true );
	}
}

void CClient::CheckTimeouts()
{
	std::vector<RFPMap::iterator> dead;

	for (RFPMap::iterator iter = m_receiving.begin(); iter != m_receiving.end(); ++iter)
	{
		if (!m_timeoutConnectionQueue.Has(iter->first))
		{
			iter->second->Timeout( m_pListener );
			dead.push_back(iter);
		}
	}

	while (!dead.empty())
	{
		m_receiving.erase( dead.back() );
		dead.pop_back();
	}
}

void CClient::HandleDataPacket( TAddr addr, BInputStream& input )
{
	RFPMap::iterator iter = m_receiving.find( addr );
	if (iter == m_receiving.end())
		SendRequestAnnouncement(addr);
	else
	{
		unsigned packet;
		input >> packet;
		size_t sz = input.Remaining();
		if (iter->second->Receive( packet, input.Get(sz), sz ))
		{
			std::cout << "Completed job " << iter->first.first << ":" << iter->first.second << std::endl;
			RFP ptr = iter->second;
			m_receiving.erase( iter );
			CClientKeepAlive keepAlive( addr, m_sock );
			if (!ptr->Complete(&keepAlive, m_pListener))
				SendRequestAnnouncement(addr);
		}
		UpdateProgressMessage();
	}
}

void CClient::HandleAnnouncePacket( TAddr addr, BInputStream& input )
{
	RFPMap::iterator iter = m_receiving.find( addr );
	if (iter == m_receiving.end())
	{
		std::string filename;
		unsigned packets;
		CHash hashSent;
		input >> filename >> packets >> hashSent;
		try
		{
			CClientKeepAlive keepAlive( addr, m_sock );
			CHash hash = GetHash(filename, &keepAlive);
			if (hash == hashSent)
			{
				std::cout << "Ignoring announce " << filename << ": already have file" << std::endl;
				return;
			}
		}
		catch (...)
		{
		}
		RFP ptr( new CReceiveFile(filename, packets, hashSent) );
		std::cout << "Starting job " << addr.first << ":" << addr.second << " (is: " << filename << ")" << std::endl;
		if (m_pListener)
			m_pListener->Notification( "Starting Job", filename );
		m_receiving.insert( std::make_pair(addr, ptr) );
		SendRequestData( addr, ptr, false, false );
	}
}

void CClient::HandleFinishPacket( TAddr addr, BInputStream& input )
{
	RFPMap::iterator iter = m_receiving.find( addr );
	if (iter == m_receiving.end())
		SendRequestAnnouncement( addr );
	else
		SendRequestData( addr, iter->second, false, true );
}

void CClient::HandleQuitPacket( TAddr addr, BInputStream& input )
{
	m_bDone = true;
}

void CClient::SendRequestAnnouncement( TAddr addr )
{
	if (m_announceRequestThrottle.Has(addr))
		return;

	BOutputStream out;
	out << 'W' << PROTOCOL_VERSION;
	m_sock.Send( addr.first, addr.second, out.GetPtr(), out.GetSize() );
	m_announceRequestThrottle.Add(1, addr);
}

void CClient::SendRequestData( TAddr addr, RFP ptr, bool islandsOnly, bool throttle )
{
	if (throttle && m_dataRequestThrottle.Has(addr))
		return;

	BOutputStream out;
	out << 'R' << PROTOCOL_VERSION;
	bool pending = ptr->GetPending( out );
	bool send = true;
	if (islandsOnly)
	{
		if (pending)
		{
			out.Rewind(sizeof(CDataRange));
			out << CDataRange(0,0);
		}
		else
			send = false;
	}
	if (send)
	{
		m_sock.Send( addr.first, addr.second, out.GetPtr(), out.GetSize() );
		m_dataRequestThrottle.Add( 1, addr );
	}
}

void CClient::UpdateProgressMessage()
{
	std::string message = "Blow: File transfer utility";

	for (RFPMap::iterator iter = m_receiving.begin(); iter != m_receiving.end(); ++iter)
		message += '\n' + iter->second->GetProgressMessage();

	if (m_pListener)
		m_pListener->Tip( message );
}

void DoClient( IClientListener * pListener )
{
	try
	{
		CClient client(pListener);
		client.Run();
	}
	catch (std::exception& e)
	{
		std::cout << "Exception in client: " << e.what() << std::endl;
	}
	catch (...)
	{
		std::cout << "Unknown exception" << std::endl;
	}

	if (pListener)
		pListener->OnQuit();
}
