#include <QApplication>

#include <QDialog>
#include <QPushbutton>
#include <QLayout>
#include <QLabel>
#include <QLineEdit>
#include <QSpinbox>
#include <QMessageBox>
#include <QFormLayout>
#include <QListView>
#include <QStringListModel>
#include <QInputDialog>
#include <QSettings>

#include <QTcpServer>
#include <QTcpSocket>
#include <QTimer>
#include <QFile>

#include <QMap>
#include <QList>
#include <QString>
#include <QDateTime>

#include "libssh2_config.h"
#include <libssh2.h>
#include <libssh2_sftp.h>

struct TunnelInfo
{
	QString m_bindAddress;
	int m_port;
	QString m_host;
	int m_hostPort;
};

class ConnectDialog: public QDialog 
{
	Q_OBJECT
public:

	ConnectDialog() : QDialog(0)
	{
		setupUi();
	}

	QString user() const { return m_user; }
	QString pass() const { return m_pass; }
	int port() const { return m_port; }
	QString host() const { return m_host; }

	QList<TunnelInfo> localTunnels() { return m_localTunnels; };

public slots:
	void setUser(const QString &user) { m_user = user; m_settings.setValue("user", user); }
	void setPass(const QString &pass) { m_pass = pass; m_settings.setValue("pass", pass); }
	void setPort(int port) { m_port = port; m_settings.setValue("port", port); }
	void setHost(const QString &host) { m_host = host; m_settings.setValue("host", host); }
	void addLocalTunnel() 
	{
		QString s = QInputDialog::getText(
			this,
			tr("Add tunnel"),
			tr("Format: [bind_address:]port:host:hostport")
		);

		QStringList l = s.split(":", QString::SkipEmptyParts);
		if(l.size() < 3 || l.size() > 4)
		{
			QMessageBox::critical(this, tr("Add tunnel"), tr("Invalid format"));
			return;
		}

		TunnelInfo lt;
		if(l.size() == 4)
		{
			lt.m_bindAddress = l.takeFirst();
		}
		lt.m_port = l.takeFirst().toInt();
		lt.m_host = l.takeFirst();
		lt.m_hostPort = l.takeFirst().toInt();
		m_localTunnels << lt;

		m_localTunnelsModel.setStringList(
			m_localTunnelsModel.stringList() << s
		);
	}
private:
	void setupUi()
	{
		QVBoxLayout *layout = new QVBoxLayout(this);
		QFormLayout *controlsLayout = new QFormLayout;

		QLineEdit *userEdit = new QLineEdit(this);
		connect(userEdit, SIGNAL(textChanged(const QString &)), 
			this, SLOT(setUser(const QString &)));
		controlsLayout->addRow(tr("Username"), userEdit); 
		userEdit->setText(m_settings.value("user").toString());

		QLineEdit *passEdit = new QLineEdit(this);
		passEdit->setEchoMode(QLineEdit::Password);
		connect(passEdit, SIGNAL(textChanged(const QString &)), 
			this, SLOT(setPass(const QString &)));
		controlsLayout->addRow(tr("Password"), passEdit);
		passEdit->setText(m_settings.value("pass").toString());

		QLineEdit *hostEdit = new QLineEdit(this);
		connect(hostEdit, SIGNAL(textChanged(const QString &)), 
			this, SLOT(setHost(const QString &)));
		controlsLayout->addRow(tr("Host"), hostEdit);
		hostEdit->setText(m_settings.value("host").toString());

		QSpinBox *portEdit = new QSpinBox(this);
		portEdit->setMinimum(1);
		portEdit->setMaximum(50000);

		connect(portEdit, SIGNAL(valueChanged(int)),
			this, SLOT(setPort(int)));
		controlsLayout->addRow(tr("Port"), portEdit);
		layout->addLayout(controlsLayout);
		portEdit->setValue(m_settings.value("port").toInt());

		QListView *listView = new QListView;
		listView->setModel(&m_localTunnelsModel);
		layout->addWidget(listView);

		QHBoxLayout *buttonLayout = new QHBoxLayout;
		QPushButton *addLocalTunnelButton = new QPushButton(tr("Add tunnel"));
		connect(addLocalTunnelButton, SIGNAL(clicked()), this, SLOT(addLocalTunnel()));
		buttonLayout->addWidget(addLocalTunnelButton);
		buttonLayout->addStretch();
		QPushButton *okButton = new QPushButton(tr("Ok"));
		connect(okButton, SIGNAL(clicked()), this, SLOT(accept()));
		okButton->setDefault(true);
		buttonLayout->addWidget(okButton);
		QPushButton *cancelButton = new QPushButton(tr("Cancel"));
		connect(cancelButton, SIGNAL(clicked()), this, SLOT(reject()));
		buttonLayout->addWidget(cancelButton);
		layout->addLayout(buttonLayout);

		resize(300, 200);
	}
private:
	QString m_user, m_pass, m_host;
	int m_port;
	QList<TunnelInfo> m_localTunnels;	
	QStringListModel m_localTunnelsModel;
	QSettings m_settings;
};

class LocalTunnelManager : public QObject
{
	Q_OBJECT
public:
	LocalTunnelManager(LIBSSH2_SESSION *session, const QList<TunnelInfo> &list) 
		:m_session(session) 
	{
		for(QList<TunnelInfo>::const_iterator it = list.begin(); it != list.end(); ++it)
		{
			setupLocalServer(*it);
		}
	}
public slots:
	void transmit()
	{

		qDebug("LocalTunnelManager: transmit() >>> %i channels", m_socketChannelMap.size());
		char buffer[4096];

		/*
		while(true)
		{
			*/

			//QApplication::processEvents();

			/*
			if(m_socketChannelMap.isEmpty())
			{
				break;
			}*/

				libssh2_session_set_blocking(m_session, 1);
			for(SocketChannelMap::iterator it = m_socketChannelMap.begin(); 
				it != m_socketChannelMap.end(); ++it)
			{

				QTcpSocket *socket = it.key();
				LIBSSH2_CHANNEL *channel = it.value();


				qDebug("LocalTunnelManager: socket->read() >>>");
				int read = socket->read(buffer, sizeof(buffer));
				QByteArray debug = QByteArray(buffer, read);
				qDebug("LocalTunnelManager: socket->read() <<< %i\n%s", read, debug.constData());
				int written = 0;
				int i = 0;
				if(read > 0)
				{
					do
					{
						qDebug("LocalTunnelManager: libssh2_channel_write() >>>");
						/*
						int i = libssh2_channel_write(
							channel, 
							buffer + written, 
							read - written
						);
						*/
						i = libssh2_channel_write(
							channel, 
							buffer, 
							read
						);
						qDebug("LocalTunnelManager: libssh2_channel_write() <<< %i", i);
						if(i < 0)
						{
							qDebug( "libssh2_channel_write: %d\n", i);
							shutdownConnection(socket, channel);
							goto next_socket;
						}
						written += i;
					} while (i > 0 && written < read);

					qDebug("LocalTunnelManager: >>> %i", read);
				}

				while(true)
				{
					qDebug("LocalTunnelManager: libssh2_channel_read() >>>");
					read = libssh2_channel_read(channel, buffer, sizeof(buffer));
					QByteArray debug = QByteArray(buffer, read);
					qDebug("LocalTunnelManager: libssh2_channel_read() <<< %i\n%s", read, debug.constData());
					if(LIBSSH2_ERROR_EAGAIN == read)
					{
						qDebug("LocalTunnelManager: LIBSSH2_ERROR_EAGAIN");
						break;
					}
					else if(read  < 0)
					{
						qDebug("LocalTunnelManager: error reading from channel");
						shutdownConnection(socket, channel);
						goto next_socket;
					}
					written = 0;
					while(written < read)
					{
						qDebug("LocalTunnelManager: socket->write() >>>");
						i = socket->write(buffer + written, read - written);
						qDebug("LocalTunnelManager: socket->write() <<< %i", i);
						if(i < 0)
						{
							qDebug("LocalTunnelManager: error writing to socket");
							shutdownConnection(socket, channel);
							goto next_socket;
						}
						written += i ;
					}

					qDebug("LocalTunnelManager: <<< %i", read);

					i = libssh2_channel_eof(channel);
					if(i) 
					{
						qDebug("LocalTunnelManager: libssh2_channel_eof %i", i);
						shutdownConnection(socket, channel);
						goto next_socket;
					}
				}
	next_socket:
				continue;


			}
				libssh2_session_set_blocking(m_session, 1);

			/*
		}
		*/


		qDebug("LocalTunnelManager: transmit() <<<");
	}
protected slots:
	void openConnection()
	{
		qDebug("LocalTunnelManager: openConnection() >>>");

		QTcpServer *server = static_cast<QTcpServer *>(sender());
		QTcpSocket *socket = server->nextPendingConnection();
		connect(socket, SIGNAL(disconnected()), this, SLOT(closeConnection()));
		if(!socket || !m_serverTunnelInfoMap.contains(server))
		{
			qDebug("LocalTunnelManager: error opening connection");
			return;
		}
		TunnelInfo tunnel = m_serverTunnelInfoMap[server];


				//libssh2_session_set_blocking(m_session, 1);
		LIBSSH2_CHANNEL *channel = libssh2_channel_direct_tcpip_ex(
			m_session, 
			tunnel.m_host.toLocal8Bit().constData(),
        	tunnel.m_hostPort, 
			server->serverAddress().toString().toLocal8Bit().constData(), 
			server->serverPort()	
		);


		if(!channel)
		{
			qDebug( "Could not open the direct-tcpip channel!\n"
					"(Note that this can be a problem at the server!"
					" Please review the server logs.)\n");
			socket->deleteLater();
			return;
		}

		
		m_socketChannelMap[socket] = channel;

		qDebug("LocalTunnelManager: openConnection() <<<");
	}
	void closeConnection()
	{
		qDebug("LocalTunnelManager: closeConnection() >>>");

		QTcpSocket *socket = static_cast<QTcpSocket *>(sender());
		LIBSSH2_CHANNEL *channel = m_socketChannelMap[socket];

		shutdownConnection(socket, channel);

		qDebug("LocalTunnelManager: closeConnection() <<<");
	}
private:
	void shutdownConnection(QTcpSocket *socket, LIBSSH2_CHANNEL *channel)
	{
		if(socket->isOpen())
		{
			socket->waitForBytesWritten();
			socket->close();
		}
		m_socketChannelMap.remove(socket);
		socket->deleteLater();

		if(channel)
		{
			/*libssh2_channel_close(channel);
			libssh2_channel_wait_closed(channel);*/
			libssh2_channel_free(channel);
		}

	}
	void setupLocalServer(const TunnelInfo &tunnel)
	{
		QTcpServer *server = new QTcpServer;
		connect(server, SIGNAL(newConnection()), this, SLOT(openConnection()));
		if(!server->listen(QHostAddress(tunnel.m_bindAddress), tunnel.m_port))
		{
			qDebug("LocalTunnelManager: fail to listen on %s:%i", 
				tunnel.m_bindAddress.toLocal8Bit().constData(), 
				tunnel.m_port);
			delete server;
		}
		qDebug("LocalTunnelManager: listening on %s:%i",
				tunnel.m_bindAddress.toLocal8Bit().constData(), 
				tunnel.m_port);
		m_serverTunnelInfoMap[server] = tunnel;
	}
private:
	LIBSSH2_SESSION *m_session;
	typedef QMap<QTcpSocket *, LIBSSH2_CHANNEL *> SocketChannelMap;
	typedef QMap<QTcpServer *, TunnelInfo> ServerTunnelInfoMap;
	SocketChannelMap m_socketChannelMap;
	ServerTunnelInfoMap m_serverTunnelInfoMap;
};


void myMessageOutput(QtMsgType type, const char *msg)
{
	static QFile *pLog = 0;
	if(!pLog)
	{
		pLog = new QFile("log.txt", QCoreApplication::instance());
		pLog->open(QIODevice::WriteOnly | QIODevice::Append | QIODevice::Text);
	}

	QString s;

	switch (type) {
	case QtDebugMsg:
		s.sprintf("Debug: %s\n", msg);
		break;
	case QtWarningMsg:
		s.sprintf("Warning: %s\n", msg);
		break;
	case QtCriticalMsg:
		s.sprintf("Critical: %s\n", msg);
		break;
	case QtFatalMsg:
		s.sprintf("Fatal: %s\n", msg);
		abort();
	}

	s = QDateTime::currentDateTime().toString(Qt::ISODate) + " " + s;
	fprintf(stderr, s.toLocal8Bit().constData());
	pLog->write(s.toLocal8Bit());
	pLog->flush();
}

int main(int argc, char *argv[])
{
	QApplication app(argc, argv);

	qInstallMsgHandler(myMessageOutput);

	QCoreApplication::setOrganizationName("Foto.ru");
	QCoreApplication::setOrganizationDomain("foto.ru");
	QCoreApplication::setApplicationName("libssh2");

	ConnectDialog dlg;
	if(dlg.exec() != QDialog::Accepted)
	{
		return -1;
	}

	int libssh2_error = libssh2_init(0);
	if(libssh2_error)
	{
		qDebug("libssh2_init() error: %d", libssh2_error);
		return -2;
	}

	QTcpSocket socket;
	socket.connectToHost(dlg.host(), dlg.port());
	if(!socket.waitForConnected())
	{
		qDebug("Error connecting to host %s", dlg.host().toLocal8Bit().constData());
		return -1;
	}

	LIBSSH2_SESSION *session = libssh2_session_init();
	if(!session)
	{
		qDebug("libssh2_session_init() failed");
		return -2;
	}

	libssh2_error = libssh2_session_startup(session, socket.socketDescriptor());
	if(libssh2_error)
	{
		qDebug("libssh2_session_startup() error: %d", libssh2_error);
		return -3;
	}

	{
    /* At this point we havn't yet authenticated.  The first thing to do
     * is check the hostkey's fingerprint against our known hosts Your app
     * may have it hard coded, may go to a file, may present it to the
     * user, that's your call
     */
    const char *fingerprint = libssh2_hostkey_hash(session, LIBSSH2_HOSTKEY_HASH_SHA1);
    qDebug( "Fingerprint: ");
    for(int i = 0; i < 20; i++)
        qDebug( "%02X ", (unsigned char)fingerprint[i]);
    qDebug( "\n");
	}

	qDebug("Password authentication: [%s] [%s]", dlg.user().toLocal8Bit().constData(), dlg.pass().toLocal8Bit().constData());
	libssh2_userauth_list(session, dlg.user().toLocal8Bit().constData(), dlg.user().toLocal8Bit().length());
	if(libssh2_userauth_password(
		session, 
		dlg.user().toLocal8Bit().constData(), 
		dlg.pass().toLocal8Bit().constData()
	))
	{
		qDebug("Password authentication failed");
		return -4;
	}

	LocalTunnelManager manager(session, dlg.localTunnels());
	QTimer timer;
	QObject::connect(&timer, SIGNAL(timeout()), &manager, SLOT(transmit()));
	timer.start(1000);
	//QTimer::singleShot(100, &manager, SLOT(transmit()));
	app.exec();

	socket.disconnectFromHost();

	libssh2_session_disconnect(session, "Client disconnecting normally");
	libssh2_session_free(session);
	libssh2_exit();

	return 0;
}

#include "qt.moc"