#include <boost/make_shared.hpp>

#include <map>

#include <iterator>

namespace Transport {

DEFINE_LOGGER(logger, "WebSocketClient");

WebSocketClient::WebSocketClient(Component *component) {

	m_component = component;

	m_upgraded = false;

#if HAVE_SWIFTEN_3

	Swift::TLSOptions o;

#endif

	m_tlsFactory = new Swift::PlatformTLSFactories();

#if HAVE_SWIFTEN_3

	m_tlsConnectionFactory = new Swift::TLSConnectionFactory(m_tlsFactory->getTLSContextFactory(), component->getNetworkFactories()->getConnectionFactory(), o);

#else

	m_tlsConnectionFactory = new Swift::TLSConnectionFactory(m_tlsFactory->getTLSContextFactory(), component->getNetworkFactories()->getConnectionFactory());

#endif

	m_reconnectTimer = m_component->getNetworkFactories()->getTimerFactory()->createTimer(1000);

	m_reconnectTimer->onTick.connect(boost::bind(&WebSocketClient::connectServer, this));

}

WebSocketClient::~WebSocketClient() {

	if (m_conn) {

		m_conn->onDataRead.disconnect(boost::bind(&WebSocketClient::handleDataRead, this, _1));

		m_conn->disconnect();

	}

	delete m_tlsFactory;

	delete m_tlsConnectionFactory;

}

void WebSocketClient::connectServer() {

	LOG4CXX_INFO(logger, "Starting DNS query for " << m_host << " " << m_path);

	m_dnsQuery = m_component->getNetworkFactories()->getDomainNameResolver()->createAddressQuery(m_host);

	m_dnsQuery->onResult.connect(boost::bind(&WebSocketClient::handleDNSResult, this, _1, _2));

	m_dnsQuery->run();

	m_reconnectTimer->stop();

}

void WebSocketClient::connectServer(const std::string &url) {

	std::string u = url.substr(6);

	m_host = u.substr(0, u.find("/"));

	m_path = u.substr(u.find("/"));

}

void WebSocketClient::write(const std::string &data) {

	if (!m_conn) {

		return;

	}

	uint8_t opcode = 129; // UTF8

	if (data.empty()) {

		LOG4CXX_INFO(logger, "pong");

		opcode = 138; // PONG

	}

	// Mask the payload

	char mask_bits[4] = {0x11, 0x22, 0x33, 0x44};

	std::string payload = data;

	for (size_t i = 0; i < data.size(); i++ ) {

		payload[i] = payload[i] ^ mask_bits[i&3];

	}

	if (data.size() <= 125) {

		uint8_t size7 = data.size() + 128; // Mask bit

		m_conn->write(Swift::createSafeByteArray(std::string((char *) &opcode, 1)

			+ std::string((char *) &size7, 1)

			+ std::string((char *) &mask_bits[0], 4)

			+ payload));

	}

	else {

		uint8_t size7 = 126 + 128; // Mask bit

		uint16_t size16 = data.size();

		size16 = htons(size16);

		m_conn->write(Swift::createSafeByteArray(std::string((char *) &opcode, 1)

			+ std::string((char *) &size7, 1)

			+ std::string((char *) &size16, 2)

			+ std::string((char *) &mask_bits[0], 4)

			+ payload));

	}

	LOG4CXX_INFO(logger, "> " << data);

}

void WebSocketClient::handleDataRead(boost::shared_ptr<Swift::SafeByteArray> data) {

	std::string d = Swift::safeByteArrayToString(*data);

	m_buffer += d;

	if (!m_upgraded) {

		if (m_buffer.find("\r\n\r\n") != std::string::npos) {

			m_buffer.erase(0, m_buffer.find("\r\n\r\n") + 4);