#include "serverconnection.h"
#include <QtCore>
#include <QPixmap>
#include <QGuiApplication>
#include <QHostInfo>
#include <unistd.h>
#include <cstdlib>
#include <sys/types.h>
#include <pwd.h>

//#define verbose
#include "../vnc/vncserver.h"

#include "../../shared/util.h"
#include "../../shared/settings.h"
#include "../util/platform/blankscreen.h"
#include "../clientapp/clientapp.h"

#define CHALLENGE_LEN 20

ServerConnection::ServerConnection(const QString& host, const quint16 port, const QByteArray& sessionName, const QByteArray& certHash, bool autoConnect) :
	QObject(NULL), _timerDelete(0), _jpegQuality(80), _authed(0), _autoConnect(autoConnect), _isLocalConnection(-1), _sessionName(sessionName), _certHash(certHash)
{
	_socket = new QSslSocket();
	_blank = new BlankScreen();
	connect(_socket, SIGNAL(encrypted()), this, SLOT(sock_connected()));
	connect(_socket, SIGNAL(readyRead()), this, SLOT(sock_dataArrival()));
	connect(_socket, SIGNAL(disconnected()), this, SLOT(sock_closed()));
	connect(_socket, SIGNAL(error(QAbstractSocket::SocketError)), this, SLOT(sock_error(QAbstractSocket::SocketError)));
	connect(_socket,
	        SIGNAL(sslErrors(const QList<QSslError> &)),
	        this,
	        SLOT(sslErrors(const QList<QSslError> &))
	       );
	connect(_socket, &QSslSocket::peerVerifyError, [=](const QSslError &error) { qDebug() << "PVE:" << error.errorString(); });
	qDebug("Connecting to %s on port %d", host.toUtf8().data(), (int)port);
	_socket->ignoreSslErrors();
	_socket->connectToHostEncrypted(host, port);
	_timerId = startTimer(4000);
	_lastData = QDateTime::currentMSecsSinceEpoch() + PING_TIMEOUT_MS;
	_timerConnectionCheck = startTimer(5000);
	// Connect the vnc start/stop signal to this class, so we can tell the server about successful vnc server startup
	connect(VncServer::instance(), SIGNAL(started(int, QString&, QString&)), this, SLOT(onVncServerStartStop(int, QString&, QString&)));
}

ServerConnection::~ServerConnection()
{
	this->disconnectFromServer();
	if (_socket != NULL) {
		qCritical("**** SOCKET DELETE IN DESTRUCTOR");
		_socket->deleteLater();
	}
	qDebug("*** Server connection destroyed.");
	_blank->deleteLater();
	_blank = NULL;
}

/**
 * Send the given message to the server.
 */
void ServerConnection::sendMessage(NetworkMessage& message)
{
	if (_socket == NULL || _socket->state() != QAbstractSocket::ConnectedState)
		return;
	message.writeMessage(_socket);
	if (!message.writeComplete()) {
		qCritical() << "SendMessage to server failed!";
	}
}

void ServerConnection::sendAttention(bool on)
{
	NetworkMessage msg;
	msg.setField(_ID, _ATTENTION);
	msg.setField(_ENABLE, _BOOL(on));
	sendMessage(msg);
}

/**
 * Disconnect from current server.
 * Do some cleanup also, like stopping any VNC server/client
 * activity, then finally close the connection.
 */
void ServerConnection::disconnectFromServer()
{
	if (_timerDelete == 0) {
		VncServer::instance()->stop();
		emit closeVnc();
		emit disconnected(this);
		_timerDelete = startTimer(500);
		qDebug("Closing connection to server");
		if (_socket != NULL) {
			_socket->blockSignals(true);
			_socket->abort();
		}
	}
}

/**
 * Handles an incoming message by the server.
 * This is somewhat of a long mess, maybe split this up some day or
 * make it OOP with a huge amount fancy features.
 */
void ServerConnection::handleMsg()
{
	_lastData = QDateTime::currentMSecsSinceEpoch() + PING_TIMEOUT_MS;
	const QString &id = _fromServer.getFieldString(_ID);

	if (_authed == 0) {
		if (id == _CHALLENGE) {
			// Initial challenge request by server
			emit stateChange(ConnectWindow::AwaitingChallengeResponse);
			_myChallenge.resize(CHALLENGE_LEN);
			for (int i = 0; i < CHALLENGE_LEN; ++i) {
				_myChallenge[i] = (char)(qrand() & 0xff);
			}
			QByteArray serverChallenge(_fromServer.getFieldBytes(_CHALLENGE));
			_toServer.reset();
			_toServer.setField(_ID, _CHALLENGE);
			_toServer.setField(_HASH, genSha1(&_sessionName, &serverChallenge));
			_toServer.setField(_CHALLENGE, _myChallenge);
			_toServer.writeMessage(_socket);
			qDebug("Received challenge, replying, sending own challenge step <- 1");
			_authed = 1;
		}
		return;
	}

	if (_authed == 1) {
		if (id == _CHALLENGE) {
			qDebug("Received challenge reply");
			if (_timerId != 0) {
				killTimer(_timerId);
				_timerId = 0;
			}
			// Check challenge response
			QByteArray serverhash(_fromServer.getFieldBytes(_HASH));
			if (serverhash != genSha1(&_sessionName, &_myChallenge)
			        && !_autoConnect) {
				qDebug("invalid. STOP.");
				emit stateChange(ConnectWindow::InvalidSslHash);
				this->disconnectFromServer();
				return;
			}
			emit stateChange(ConnectWindow::LoggingIn);
			const char *user = getpwuid(getuid())->pw_name;
			if (user == NULL || *user == '\0')
				user = getenv("USER");
			if (user == NULL || *user == '\0')
				user = getenv("USERNAME");
			if (user == NULL || *user == '\0')
				user = "Hans Affe";
			_toServer.reset();
			_toServer.setField(_ID, _LOGIN);
			_toServer.setField("HOST", QHostInfo::localHostName());
			_toServer.setField("NAME", QString(user));
			qDebug() << "logging into manager, exam mode is " << clientApp->isExamMode();
			_toServer.setField(_EXAMMODE, clientApp->isExamMode() ? __TRUE : __FALSE);
			/* TODO: (Question) Why is this here not using sendMessage() ? */
			qDebug("Sending login request!");
			if (_toServer.writeMessage(_socket)) {
				_authed = 2;
				qDebug("valid, step <- 2");
			} else {
				this->disconnectFromServer();
			}
		}
		return;
	}

	if (_authed == 2) {
		if (id == _LOGIN) {
			qDebug("login accepted, step <- 3");
			_authed = 3;
			emit stateChange(ConnectWindow::Connected);
		}
		return;
	}

	// message THUMB - server requests screenshot as thumbnail
	if (id == _THUMB) {
		if (clientApp->isExamMode()) {
			QByteArray emptyArray;
			_toServer.setField(_ID, _THUMB);
			_toServer.setField(_IMG, emptyArray);
			sendMessage(_toServer);
			return;
		}
		int x = _fromServer.getFieldString(_X).toInt();
		int y = _fromServer.getFieldString(_Y).toInt();

		QRect primaryRect = QGuiApplication::primaryScreen()->geometry();
		// Limit requested size so it won't easily leak sensitive information
		if (x < 32) {
			x = 32;
		} else if (x > primaryRect.width() / 8) {
			x = primaryRect.width() / 8;
		}
		if (y < 18) {
			y = 18;
		} else if (y > primaryRect.height() / 8) {
			y = primaryRect.height() / 8;
		}

		QPixmap desktop(
			QGuiApplication::primaryScreen()->grabWindow(0,
				primaryRect.x(), primaryRect.y(), primaryRect.width(), primaryRect.height()
			).scaled(x, y, Qt::KeepAspectRatio, Qt::SmoothTransformation
		));
		QByteArray bytes;
		QBuffer jpgBuffer(&bytes);
		jpgBuffer.open(QIODevice::WriteOnly);
		if (desktop.save(&jpgBuffer, "JPG", _jpegQuality)) { // writes pixmap into bytes in JPG format
			// Try to adjust quality so we stay between 3 and 4.5 KB
			if (_jpegQuality < 90 && bytes.size() < 3000)
				_jpegQuality += 7;
			else if (_jpegQuality > 40 && bytes.size() > 4500)
				_jpegQuality -= 7;
		} else {
			// FALLBACK
			bytes.clear();
			QBuffer pngBuffer(&bytes);
			pngBuffer.open(QIODevice::WriteOnly);
			if (!desktop.save(&pngBuffer, "PNG")) { // writes pixmap into bytes in PNG format
				qDebug("Could not convert screenshot to PNG nor JPG");
				return; // Failed :-(
			}
		}
		_toServer.reset();
		_toServer.setField(_ID, _THUMB);
		_toServer.setField(_IMG, bytes);
		sendMessage(_toServer);
	} // message VNCSERVER - start local vncserver
	else if (id == _VNCSERVER) {
		if (clientApp->isExamMode()) {
			qDebug() << "denied request for vnc server (exam mode)";
			return;
		}
		const bool enable = (_fromServer.getFieldString("ENABLE").toInt() != 0);
		if (enable) {
			emit closeVnc(); // In case we are watching some other client, stop doing so
			VncServer::instance()->start();
		} else {
			VncServer::instance()->stop();
		}
	} else if (id == _VNCCLIENT) {
		if (clientApp->isExamMode()) {
			qDebug() << "denied request for vnc projection (exam mode)";
			return;
		}
		const QString host(_fromServer.getFieldString("HOST"));
		const int port = _fromServer.getFieldString("PORT").toInt();
		if (host.isEmpty() || port <= 0) {
			emit closeVnc();
		} else {
			emit openVnc(host, port, _fromServer.getFieldString("ROPASS"), true, true, _fromServer.getFieldString("CAPTION"), _fromServer.getFieldString("CLIENTID").toInt(), _fromServer.getFieldBytes(_THUMB));
		}
	} else if (id == _LOCK) {
		const bool enable = (_fromServer.getFieldString("ENABLE").toInt() != 0);
		if (enable && !clientApp->isConnectedToLocalManager()) {
			_blank->lock(_fromServer.getFieldString("MESSAGE"));
		} else {
			_blank->unlock();
		}
	} else if (id == _ATTENTION) {
		emit attentionChanged(_fromServer.getFieldString(_ENABLE) == __TRUE);
	}
}

void ServerConnection::checkLocalConnection()
{
	if (_socket == NULL) {
		return;
	}
	if (_socket->peerAddress() == _socket->localAddress()) {
		_isLocalConnection = 1;
	} else {
		_isLocalConnection = 0;
	}
}

/*
 * Override
 */

void ServerConnection::timerEvent(QTimerEvent *event)
{
	if (event->timerId() == _timerConnectionCheck) {
		if (_lastData < QDateTime::currentMSecsSinceEpoch()) {
			qDebug() << "Ping timeout";
			this->disconnectFromServer();
			killTimer(_timerConnectionCheck);
		}
	} else if (event->timerId() == _timerId) {
		qDebug() << "Connect timeout";
		killTimer(_timerId);
		_timerId = 0;
		this->disconnectFromServer();
	} else if (event->timerId() == _timerDelete) {
		if (_socket == NULL || _socket->state() == QAbstractSocket::UnconnectedState) {
			if (_socket != NULL)
				_socket->deleteLater();
			_socket = NULL;
			killTimer(_timerDelete);
			this->deleteLater();
			return;
		}
		_socket->abort();
		qDebug("A socket is still pending...");
	} else
		killTimer(event->timerId());
}

/*
 * Slots
 */

/**
 * This slot is triggered by the vnc server runner once the external VNC
 * server was succesfully started, or was terminated (either planned or
 * crashed).
 */
void ServerConnection::onVncServerStartStop(int port, QString& ropass, QString& rwpass)
{
	_toServer.reset();
	_toServer.setField(_ID, _VNCSERVER);
	if (port <= 0) {
		_toServer.setField("PORT", QByteArray("0"));
	} else {
		_toServer.setField("PORT", QString::number(port));
		_toServer.setField("ROPASS", ropass);
		_toServer.setField("RWPASS", rwpass);
	}
	sendMessage(_toServer);
}

/**
 * This slot is triggered once the internal VNC viewer has started or stopped
 * displaying a VNC stream. We'll inform the server about the state change.
 */
void ServerConnection::onVncViewerStartStop(const bool started, const int clientId)
{
	_toServer.reset();
	_toServer.setField(_ID, _VNCCLIENT);
	if (started)
		_toServer.setField("ENABLED", __TRUE);
	else
		_toServer.setField("ENABLED", __FALSE);
	_toServer.setField("CLIENTID", QString::number(clientId));
	sendMessage(_toServer);
}

/**
 * An ssl error happened. If it's an expected one, we ignore it
 * and keep going. Otherwise the connection will be terminated.
 */
void ServerConnection::sslErrors(const QList<QSslError> & errors)
{
	_socket->ignoreSslErrors();
	for (QList<QSslError>::const_iterator it = errors.begin(); it != errors.end(); it++) {
		const QSslError &err = *it;
		qDebug("Connect SSL: %s", qPrintable(err.errorString()));
		if (err.error() == QSslError::HostNameMismatch)
			continue; // We don't pay attention to hostnames for validation
		if (err.error() == QSslError::SelfSignedCertificate)
			continue; // Also, this will always be the case; we check the fingerprint later
		if (err.error() == QSslError::CertificateNotYetValid || err.error() == QSslError::CertificateExpired)
			continue;
		// Some other SSL error - do not ignore
		qDebug() << "FATAL!";
		_socket->ignoreSslErrors(QList<QSslError>());
		return;
	}
}

void ServerConnection::sock_dataArrival()
{
	if (_socket == NULL || _socket->state() != QAbstractSocket::ConnectedState) {
		qDebug("dataArrival called in bad state");
		return;
	}

	while (_socket->bytesAvailable() > 0) {
		int retval = _fromServer.readMessage(_socket); // let the message read data from socket
		if (retval == NM_READ_FAILED) { // error parsing msg, disconnect client!
			this->disconnectFromServer();
			return;
		}
		if (retval == NM_READ_INCOMPLETE)
			return;
		if (_fromServer.readComplete()) { // message is complete
			this->handleMsg();
			if (_socket == NULL)
				return;
			_fromServer.reset();
		}
	}
}

void ServerConnection::sock_closed()
{
	// should this be unreliable in some way i suggest using the signal "stateChanged()" instead
	// and check if the state changed to unconnected.
	qDebug("Socket was closed... oh well..");
	this->disconnectFromServer();
}

void ServerConnection::sock_error(QAbstractSocket::SocketError errcode)
{
	qDebug("Connection error: %d", (int)errcode);
	this->disconnectFromServer();
}

void ServerConnection::sock_connected()
{
	qDebug() << "Connection to server established and encrypted";
	QByteArray cert(_socket->peerCertificate().digest(QCryptographicHash::Sha1));
	if (_certHash != cert) {
		emit stateChange(ConnectWindow::InvalidCert);
		this->disconnectFromServer();
		return;
	}
	emit stateChange(ConnectWindow::AwaitingChallenge);
}