#include "clientsocket.h"
#include "clientauthenticationresponsemessage.h"
#include "clientinitmessage.h"
#include "clientkeyeventmessage.h"
#include "clientnativekeyeventmessage.h"
#include "clientmouseeventmessage.h"
#include "clientrequestframebufferupdatemessage.h"
#include "clientsetencodingsmessage.h"
#include "clientsetpixelformatmessage.h"
#include "serverauthenticationchallengemessage.h"
#include "serverframebufferinitmessage.h"
#include "serverframebufferupdatemessage.h"
#include "serverprotocolversionmessage.h"
#include "serversecurity33message.h"
#include "serversecurity37message.h"
#include "serversecurityresultmessage.h"
#include "serversecurity37resultmessage.h"
#include "clientversionmessage.h"
#include "clientsecuritytypemessage.h"
#include "compressioncontext.h"
#include "clientconfigureaudiomessage.h"
#include "serveraudiocapturenotifymessage.h"
#include "servervencryptauthtyperesponsemessage.h"
#include "servervencryptversionmessage.h"
#include "clientvencryptversionmessage.h"
#include "clientvencryptauthtypesmessage.h"
#include "clientvencryptplainauthenticationmessage.h"
#include "servervencryptauthresponse.h"
#include "connectionprofiler.h"
#include "rawrectreader.h"

#include <QReadLocker>
#include <QSslSocket>
#include <QHostAddress>
#include <QTimer>
#include <QGLWidget>
#include <QDebug>
#include <QFinalState>
#include <QAudioOutput>

void ClientSocket::configureStates()
{
    QState *disconnected = new QState(m_stateMachine);
    QState *handshaking = new QState(m_stateMachine);
    QState *connected = new QState(m_stateMachine);
    QFinalState *sessionDone = new QFinalState(m_stateMachine);

    QState *idle = new QState(disconnected);
    QState *socketConnected = new QState(disconnected);
    //QState *socketError = new QState(disconnected);

    QState *handshakeReadVersion = new QState(handshaking);
    QState *handshakeSendVersionResponse = new QState(handshaking);
    QState *handshakeVersionError = new QState(handshaking);
    QState *handshakeReadSecurityTypes = new QState(handshaking);
    QState *handshakeSendSecurityType = new QState(handshaking);
    QState *handshakeReadSecurityAuthResult = new QState(handshaking);
    QState *handshakeSecurityError = new QState(handshaking);
    QState *handshakeReadAuthChallenge = new QState(handshaking);
    QState *handshakeReadVencryptVersion = new QState(handshaking);
    QState *handshakeReadVencryptSubauth = new QState(handshaking);
    QState *handshakeDoVencryptAuth = new QState(handshaking);
    QState *handshakeSendVencryptPlainAuth = new QState(handshaking);
    QState *handshakeStartTLS = new QState(handshaking);
    QState *handshakeStartTLSVNC = new QState(handshaking);
    QState *handshakeStartTLSPlain = new QState(handshaking);
    QState *handshakeStartX509 = new QState(handshaking);
    QState *handshakeStartX509VNC = new QState(handshaking);
    QState *handshakeStartX509Plain = new QState(handshaking);
    QState *handshakeReadVencryptAuthResult = new QState(handshaking);
    QState *handshakeSendVNCAuth = new QState(handshaking);
    QState *handshakeNoAuth38 = new QState(handshaking);
    QState *handshakeSendClientInit = new QState(handshaking);
    QState *handshakeReadFramebufferSetup = new QState(handshaking);
    QState *handshakeFramebufferError = new QState(handshaking);

    QState *handleFramebufferUpdateState = new QState(connected);
    QState *readFrameBufferMessageTypeState = new QState(handleFramebufferUpdateState);
    QState *readQemuUpdateMessageState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateMessageState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateParseNextMessageState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectDesktopResizeState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectWMViState =  new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectPointerTypeChangeState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectAudioState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectNativeKeysState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectRawState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectRREState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectCoRREState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectCopyRectState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectHextileState = new QState(handleFramebufferUpdateState);
    QState *readRfbUpdateRectZlibState = new QState(handleFramebufferUpdateState);
    QState *readQemuUpdateAudioMessageState = new QState(handleFramebufferUpdateState);
    m_rawRectReadState = new NativeRawRectReader( m_connection, readRfbUpdateRectRawState );

    QFinalState *handleFramebufferUpdateDoneState = new QFinalState(handleFramebufferUpdateState);

    disconnected->addTransition( m_connection, SIGNAL(connected()), socketConnected );
    //disconnected->addTransition( m_connection, SIGNAL(error(QAbstractSocket::SocketError)), socketError );

    //socketError->addTransition(sessionDone);

    socketConnected->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadVersion );

    handshaking->addTransition( m_connection, SIGNAL(disconnected()), sessionDone );
    //handshaking->addTransition( m_connection, SIGNAL(error(QAbstractSocket::SocketError)), socketError );

    handshakeReadVersion->addTransition(this, SIGNAL(versionError()), handshakeVersionError);
    handshakeReadVersion->addTransition(this, SIGNAL(versionSupported()), handshakeSendVersionResponse );

    handshakeSendVersionResponse->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadSecurityTypes);
    handshakeReadSecurityTypes->addTransition(this, SIGNAL(securityTypeSupported()), handshakeSendSecurityType );
    handshakeReadSecurityTypes->addTransition(this, SIGNAL(securityTypeNotSupported()), handshakeSecurityError );

    handshakeSendSecurityType->addTransition(this, SIGNAL(useVNCAuth()), handshakeReadAuthChallenge );
    handshakeSendSecurityType->addTransition(this, SIGNAL(skipResponse()), handshakeSendClientInit);
    handshakeSendSecurityType->addTransition(this, SIGNAL(useNoAuth()), handshakeNoAuth38);
    handshakeSendSecurityType->addTransition(this, SIGNAL(useVencryptAuth()), handshakeDoVencryptAuth);
    handshakeSendSecurityType->addTransition(this, SIGNAL(securityTypeNotSupported()), handshakeSecurityError);

    handshakeReadAuthChallenge->addTransition(m_connection, SIGNAL(readyRead()), handshakeSendVNCAuth );

    handshakeSendVNCAuth->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadSecurityAuthResult );
    handshakeNoAuth38->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadSecurityAuthResult );

    handshakeReadSecurityAuthResult->addTransition(this, SIGNAL(authSuccessful()), handshakeSendClientInit );
    handshakeReadSecurityAuthResult->addTransition(this, SIGNAL(authError()), handshakeSecurityError );

    handshakeDoVencryptAuth->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadVencryptVersion );
    handshakeReadVencryptVersion->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadVencryptSubauth);

    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(authError()), handshakeSecurityError );
    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(usePlainAuth()), handshakeSendVencryptPlainAuth );

    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useTLS()), handshakeStartTLS );
    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useTLSVNC()), handshakeStartTLSVNC );
    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useTLSPlain()), handshakeStartTLSPlain );

    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useX509()), handshakeStartX509 );
    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useX509VNC()), handshakeStartX509VNC );
    handshakeReadVencryptSubauth->addTransition(this, SIGNAL(useX509Plain()), handshakeStartX509Plain );

    handshakeStartTLS->addTransition(m_connection, SIGNAL(encrypted()), handshakeReadVencryptAuthResult );
    handshakeStartX509->addTransition(m_connection, SIGNAL(encrypted()), handshakeReadVencryptAuthResult );

    handshakeStartTLSVNC->addTransition(m_connection, SIGNAL(encrypted()), handshakeReadAuthChallenge );
    handshakeStartX509VNC->addTransition(m_connection, SIGNAL(encrypted()), handshakeReadAuthChallenge );

    handshakeStartTLSPlain->addTransition(m_connection, SIGNAL(encrypted()), handshakeSendVencryptPlainAuth );
    handshakeStartX509Plain->addTransition(m_connection, SIGNAL(encrypted()), handshakeSendVencryptPlainAuth );

    handshakeSendVencryptPlainAuth->addTransition( m_connection, SIGNAL(readyRead()), handshakeReadVencryptAuthResult );

    handshakeReadVencryptAuthResult->addTransition(this, SIGNAL(authSuccessful()), handshakeSendClientInit );
    handshakeReadVencryptAuthResult->addTransition(this, SIGNAL(authError()), handshakeSecurityError );

    handshakeSendClientInit->addTransition(m_connection, SIGNAL(readyRead()), handshakeReadFramebufferSetup );

    handshakeReadFramebufferSetup->addTransition( this, SIGNAL(framebufferError()), handshakeFramebufferError );
    handshakeReadFramebufferSetup->addTransition( m_connection, SIGNAL(readyRead()), handleFramebufferUpdateState);

    handshakeSecurityError->addTransition(sessionDone);
    handshakeFramebufferError->addTransition(sessionDone);

    connected->addTransition( m_connection, SIGNAL(disconnected()), sessionDone );
    //connected->addTransition( m_connection, SIGNAL(error(QAbstractSocket::SocketError)), socketError );
    connected->addTransition( m_connection, SIGNAL(readyRead()), handleFramebufferUpdateState );
    connected->addTransition( this, SIGNAL(moreBytesAvailable()), handleFramebufferUpdateState );

    readFrameBufferMessageTypeState->addTransition(this, SIGNAL(gotQemuUpdateMessage()), readQemuUpdateMessageState);
    readFrameBufferMessageTypeState->addTransition(this, SIGNAL(gotRfbUpdateMessage()), readRfbUpdateMessageState);
    readFrameBufferMessageTypeState->addTransition(this, SIGNAL(finishedReadingMessage()), handleFramebufferUpdateDoneState);

    readRfbUpdateMessageState->addTransition(this, SIGNAL(parseNextRfbUpdateRect()), readRfbUpdateParseNextMessageState);
    readRfbUpdateMessageState->addTransition(this, SIGNAL(finishedReadingMessage()), handleFramebufferUpdateDoneState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotDesktopResizeMessage()), readRfbUpdateRectDesktopResizeState );
    readRfbUpdateRectDesktopResizeState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotWMViMessage()), readRfbUpdateRectWMViState);
    readRfbUpdateRectWMViState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotPointerTypeChangeMessage()), readRfbUpdateRectPointerTypeChangeState);
    readRfbUpdateRectPointerTypeChangeState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotAudioMessage()), readRfbUpdateRectAudioState);
    readRfbUpdateRectAudioState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotNativeKeysMessage()), readRfbUpdateRectNativeKeysState);
    readRfbUpdateRectNativeKeysState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectRawMessage()), readRfbUpdateRectRawState);
    readRfbUpdateRectRawState->addTransition(m_rawRectReadState);
    m_rawRectReadState->addTransition(readRfbUpdateParseNextMessageState);


    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectRREMessage()), readRfbUpdateRectRREState);
    readRfbUpdateRectRREState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectCoRREMessage()), readRfbUpdateRectCoRREState);
    readRfbUpdateRectCoRREState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectCopyRectMessage()), readRfbUpdateRectCopyRectState);
    readRfbUpdateRectCopyRectState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectHextileMessage()), readRfbUpdateRectHextileState);
    readRfbUpdateRectHextileState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(gotRectZlibMessage()), readRfbUpdateRectZlibState);
    readRfbUpdateRectZlibState->addTransition(readRfbUpdateParseNextMessageState);

    readRfbUpdateParseNextMessageState->addTransition(this, SIGNAL(finishedReadingMessage()), handleFramebufferUpdateDoneState);


    readQemuUpdateMessageState->addTransition(this, SIGNAL(gotAudioMessage()), readQemuUpdateAudioMessageState );
    readQemuUpdateAudioMessageState->addTransition(this, SIGNAL(finishedReadingMessage()), handleFramebufferUpdateDoneState );


    connect( disconnected, SIGNAL(entered()), this, SLOT(onDisconnected()));
    connect( handshaking, SIGNAL(entered()), this, SLOT(onHandshaking()));
    connect( connected, SIGNAL(entered()), this, SLOT(onConnected()));
    connect( sessionDone, SIGNAL(entered()), this, SLOT(onSessionDone()));
    connect( socketConnected, SIGNAL(entered()), this, SLOT(onSocketConnected()));
    //connect( socketError, SIGNAL(entered()), this, SLOT(onSocketError()));
    connect( handshakeReadVersion, SIGNAL(entered()), this, SLOT(onHandshakeReadVersion()));
    connect( handshakeSendVersionResponse, SIGNAL(entered()), this, SLOT(onHandshakeSendVersionResponse()));
    connect( handshakeVersionError, SIGNAL(entered()), this, SLOT(onHandshakeVersionError()));
    connect( handshakeReadSecurityTypes, SIGNAL(entered()), this, SLOT(onHandshakeReadSecurityTypes()));
    connect( handshakeSendSecurityType, SIGNAL(entered()), this, SLOT(onHandshakeSendSecurityType()));
    connect( handshakeReadSecurityAuthResult, SIGNAL(entered()), this, SLOT(onHandshakeReadSecurityAuthResult()));
    connect( handshakeSecurityError, SIGNAL(entered()), this, SLOT(onHandshakeSecurityError()));
    connect( handshakeSendVNCAuth, SIGNAL(entered()), this, SLOT(onHandshakeSendVNCAuth()));
    connect( handshakeReadVencryptSubauth, SIGNAL(entered()), this, SLOT(onHandshakeReadVencryptSubauth()));

    connect( handshakeDoVencryptAuth, SIGNAL(entered()), this, SLOT(onHandshakeDoVencryptAuth()));
    connect( handshakeReadVencryptVersion, SIGNAL(entered()), this, SLOT(onHandshakeReadVencryptVersion()));
    connect( handshakeSendVencryptPlainAuth, SIGNAL(entered()), this, SLOT(onHandshakeSendVencryptPlainAuth()));
    connect( handshakeStartTLS, SIGNAL(entered()), this, SLOT(onHandshakeStartTLS()));
    connect( handshakeStartTLSVNC, SIGNAL(entered()), this, SLOT(onHandshakeStartTLSVNC()));
    connect( handshakeStartTLSPlain, SIGNAL(entered()), this, SLOT(onHandshakeStartTLSPlain()));
    connect( handshakeStartX509, SIGNAL(entered()), this, SLOT(onHandshakeStartX509()));
    connect( handshakeStartX509VNC, SIGNAL(entered()), this, SLOT(onHandshakeStartX509VNC()));
    connect( handshakeStartX509Plain, SIGNAL(entered()), this, SLOT(onHandshakeStartX509Plain()));
    connect( handshakeReadVencryptAuthResult, SIGNAL(entered()), this, SLOT(onHandshakeReadVencryptAuthResult()));

    connect( handshakeSendClientInit, SIGNAL(entered()), this, SLOT(onHandshakeSendClientInit()));
    connect( handshakeReadFramebufferSetup, SIGNAL(entered()), this, SLOT(onHandshakeReadFramebufferSetup()));
    connect( handshakeFramebufferError, SIGNAL(entered()), this, SLOT(onHandshakeFramebufferError()));
    connect( handleFramebufferUpdateState, SIGNAL(entered()), this, SLOT(onHandleFramebufferUpdateState()) );
    connect( readRfbUpdateParseNextMessageState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateParseNextMessageState()) );
    connect( readFrameBufferMessageTypeState, SIGNAL(entered()), this, SLOT(onReadFrameBufferMessageTypeState()) );
    connect( readQemuUpdateMessageState, SIGNAL(entered()), this, SLOT(onReadQemuUpdateMessageState()) );
    connect( readRfbUpdateMessageState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateMessageState()) );
    connect( readRfbUpdateRectDesktopResizeState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectDesktopResizeState()) );
    connect( readRfbUpdateRectWMViState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectWMViState()) );
    connect( readRfbUpdateRectPointerTypeChangeState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectPointerTypeChangeState()) );
    connect( readRfbUpdateRectAudioState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectAudioState()) );
    connect( readRfbUpdateRectNativeKeysState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectNativeKeysState()) );
    connect( readRfbUpdateRectRawState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectRawState()) );
    connect( readRfbUpdateRectRREState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectRREState()) );
    connect( readRfbUpdateRectCoRREState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectCoRREState()) );
    connect( readRfbUpdateRectCopyRectState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectCopyRectState()) );
    connect( readRfbUpdateRectHextileState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectHextileState()) );
    connect( readRfbUpdateRectZlibState, SIGNAL(entered()), this, SLOT(onReadRfbUpdateRectZlibState()) );
    connect( readQemuUpdateAudioMessageState, SIGNAL(entered()), this, SLOT(onReadQemuUpdateAudioMessageState()) );
    connect( handleFramebufferUpdateDoneState, SIGNAL(entered()), this, SLOT(onHandleFramebufferUpdateDoneState()) );

    readRfbUpdateRectRawState->setInitialState(m_rawRectReadState);
    handleFramebufferUpdateState->setInitialState(readFrameBufferMessageTypeState);
    disconnected->setInitialState(idle);
    m_stateMachine->setInitialState(disconnected);

}

ClientSocket::ClientSocket( QObject *parent ) : QObject(parent),
    m_authType(ClientGlobals::None),
    m_subAuthType(ClientGlobals::TLSNone),
    m_authenticated(false),
    m_maxClientVersion(ClientGlobals::RFB_38),
    m_shared(false),
    m_canAuthenticate(false),
    m_audioSupported(false),
    m_fbBits(0),
    m_useNativeKeys(false),
    m_enableAudio(false),
    m_audioDevice(0),
    m_audioBuffer(0)
{
    m_stateMachine = new QStateMachine(this);
    m_unzip = new CompressionContext();
    m_connection = new QSslSocket(this);

    connect( m_connection, SIGNAL(sslErrors(QList<QSslError>)), this, SLOT(handleSslError(QList<QSslError>)));
    connect( m_connection, SIGNAL(error(QAbstractSocket::SocketError)), this, SLOT(handleError(QAbstractSocket::SocketError)));

    configureStates();
    m_stateMachine->start();
}

ClientSocket::~ClientSocket()
{
    delete m_unzip;
}

void ClientSocket::setAudioDevice( QAudioOutput *device )
{
    m_audioDevice = device;
}

void ClientSocket::connectToHost( const QHostAddress &hostName, quint16 port )
{
    qDebug() << "Connecting to" << hostName << port;
    m_connection->connectToHost(hostName, port);

}
void ClientSocket::disconnectFromHost( )
{
    m_connection->disconnectFromHost();
    qDebug() << "Disconnect";
    emit disconnected();
}

void ClientSocket::sendKeyEvent( const KeyEvent &event )
{
    ClientKeyEventMessage message;
    message.setKeyEvent(event);
    //qDebug() << "Send xkeys keys" << message << m_connection->bytesToWrite();
    message.writeBytes(m_connection);
}

void ClientSocket::sendNativeKeyEvent( const KeyEvent &event )
{
    ClientNativeKeyEventMessage message;
    message.setKeyEvent(event);
    //qDebug() << "Send native keys" << message << m_connection->bytesToWrite();
    message.writeBytes(m_connection);
}

void ClientSocket::sendMouseEvent( const MouseEvent &event )
{
    ClientMouseEventMessage message;
    message.setMouseEvent(event);
    message.writeBytes(m_connection);
    //qDebug() << "Send mouse event" << m_connection->bytesToWrite();
}

void ClientSocket::sendUpdateRect( const QRect &rect )
{
    ClientRequestFramebufferUpdateMessage message;
    message.setUpdateRect(rect);
    message.writeBytes(m_connection);
}

void ClientSocket::handleConnect( )
{
    qDebug() << "Connected...";
}

void ClientSocket::handleDisconnect( )
{
     qDebug() << "Disconnected...";
     emit disconnected();
}

void ClientSocket::handleError( const QAbstractSocket::SocketError &err )
{
    switch( err )
    {
        case QAbstractSocket::AddressInUseError:
            emit error( tr("Address in use") );
        break;
        default:
            emit error( tr("Generic error.") );
    }
}

void ClientSocket::sendClientVersion()
{
    ClientVersionMessage message;
    message.setMaxVersion( m_maxClientVersion );
    message.writeBytes(m_connection);
}

void ClientSocket::sendSecurityType()
{

    ClientSecurityTypeMessage message;
    message.setSecurityType( m_authType );
    qDebug() << "Send security type: " << message;

    message.writeBytes(m_connection);
}

void ClientSocket::sendPassword()
{
    ClientAuthenticationResponseMessage message;
    message.setPassword(m_password);
    message.setChallenge(m_challenge);
    message.writeBytes(m_connection);
}

void ClientSocket::sendClientInit()
{
    ClientInitMessage message;
    message.setShared(m_shared);
    message.writeBytes(m_connection);
}

void ClientSocket::sendPixelFormat()
{
    ClientSetPixelFormatMessage message;
    message.setPixelFormat(m_pixelFormat);
    message.writeBytes( m_connection );
}

void ClientSocket::setUseNativeKeys( bool useNativeKeys )
{
    m_useNativeKeys = useNativeKeys;
}

void ClientSocket::setEncodingPriorites( const QStringList &encodings )
{
    m_encodingPriorites = encodings;
}

void ClientSocket::setCaCertPath( const QString &path )
{
    m_caCertPath = path;
}

void ClientSocket::setAuthSubtype( const ClientGlobals::SubAuthType &type )
{
    m_subAuthType = type;
}

void ClientSocket::setAudioEnabled( bool enabled )
{
    m_enableAudio = enabled;
}

void ClientSocket::sendEncodings()
{
    ClientSetEncodingsMessage message;
    QList<ClientGlobals::RfbEncoding> encodings;
    encodings << ClientGlobals::PointerTypeChange
              << ClientGlobals::CopyRect
              << ClientGlobals::RRE
              << ClientGlobals::CoRRE
              << ClientGlobals::DesktopResize
              << ClientGlobals::WMVi;

    if ( m_enableAudio )
        encodings << ClientGlobals::Audio;

    if ( m_useNativeKeys )
        encodings << ClientGlobals::NativeKeys;

    foreach( QString enc, m_encodingPriorites )
        encodings << ClientGlobals::mapRfbEncoding( enc );

    if( m_encodingPriorites.isEmpty() )
        encodings << ClientGlobals::Raw;

    message.setEncodings( encodings );
    message.writeBytes(m_connection);
}

void ClientSocket::readVersion()
{
    ServerProtocolVersionMessage message = ServerMessageBase::readMessage<ServerProtocolVersionMessage>(m_connection);
    qDebug() << message;
    m_maxClientVersion = message.version();
    // check to see is it is at least 3.7, or we bail.
    emit versionSupported();
}

void ClientSocket::readSecurityTypes()
{
    ServerSecurity37Message message;
    if( m_maxClientVersion == ClientGlobals::RFB_33 )
    {
        message = ServerMessageBase::readMessage<ServerSecurity33Message>(m_connection);
        qDebug() << "Oldstyle security:" << message;
    }
    else
    {
        message = ServerMessageBase::readMessage<ServerSecurity37Message>(m_connection);
        qDebug() << "Newstyle security:" << message;
    }

    if( message.securityTypes().count() > 0 )
    {
        if( message.securityTypes().contains(ClientGlobals::AuthInvalid) )
            m_canAuthenticate = false;
        else
            m_canAuthenticate = ( message.securityTypes().contains( m_authType ) );
    }
    else
    {
        m_canAuthenticate = false;
    }

}

void ClientSocket::readSecurityResult()
{    
    if( m_maxClientVersion == ClientGlobals::RFB_38 )
    {
        ServerSecurityResultMessage message = ServerMessageBase::readMessage<ServerSecurityResultMessage>(m_connection);
        qDebug() << "Newstyle security:" << message;
        m_canAuthenticate = ( !message.hasError() );

    }
    else
    {
        ServerSecurity37ResultMessage message = ServerMessageBase::readMessage<ServerSecurity37ResultMessage>(m_connection);
        qDebug() << "Oldstyle security:" << message;
        m_canAuthenticate = ( !message.hasError() );
    }
}

void ClientSocket::readChallenge()
{
    ServerAuthenticationChallengeMessage message = ServerMessageBase::readMessage<ServerAuthenticationChallengeMessage>(m_connection);
    m_challenge = message.challenge();
}

void ClientSocket::readAuthResult()
{
    ServerSecurityResultMessage message = ServerMessageBase::readMessage<ServerSecurityResultMessage>(m_connection);
    m_authenticated = !message.hasError();
    qDebug() << "Auth result" << message;
}

void ClientSocket::updateFrambufferFromPixelFormat( const QSize &size )
{
    if ( m_pixelFormat.bpp() == 32 )
        m_framebuffer = QImage( size, QImage::Format_RGB32);
    else if ( m_pixelFormat.bpp() == 16 )
        m_framebuffer = QImage( size, QImage::Format_RGB16);
    else
        m_framebuffer = QImage( size, QImage::Format_ARGB32_Premultiplied);

    m_framebuffer.fill(0);
    m_fbBits = m_framebuffer.scanLine(0);
}

void ClientSocket::readInitFramebuffer()
{
    ServerFramebufferInitMessage message = ServerMessageBase::readMessage<ServerFramebufferInitMessage>(m_connection);
    qDebug() << message;
    m_pixelFormat = message.pixelFormat();
    m_connectionName = message.desktopName();
    emit framebufferResize( message.size() );
    emit desktopNameChanged(m_connectionName);
    updateFrambufferFromPixelFormat(message.size());
}

void ClientSocket::setAuthenticationType( const ClientGlobals::AuthType &type )
{
    m_authType = type;
}

void ClientSocket::setPassword( const QByteArray &pass )
{
    m_password = pass;
}

void ClientSocket::setUsername( const QByteArray &username )
{
    m_username = username;
}

void ClientSocket::setSharedSession( bool share )
{
    m_shared = share;
}

void ClientSocket::sendEnableAudio( )
{
    if ( m_audioSupported && m_audioDevice)
    {
        quint8 format = ClientGlobals::mapQtAudioFormatToRFB( m_audioDevice->format() );
        if( format != 0x06 )
        {
            qDebug("Enable audio");
            ClientConfigureAudioMessage configureMessage;
            configureMessage.setAudioFormat( m_audioDevice->format() );

            ClientAddAudioMessage addMessage;

            configureMessage.writeBytes(m_connection);
            addMessage.writeBytes(m_connection);
        }
    }
}

void ClientSocket::sendDisableAudio( )
{
    qDebug("Disable audio");
    ClientRemoveAudioMessage removeMessage;
    removeMessage.writeBytes(m_connection);
}

void ClientSocket::readVencryptVersion()
{
    ServerVencryptVersionMessage message = ServerMessageBase::readMessage<ServerVencryptVersionMessage>(m_connection);
    if( message.majorVersion() != 0 ||
        message.minorVersion() != 2 )
    {
        disconnectFromHost();
    }
    else
    {
        ClientVencryptVersionMessage sendMessage;
        sendMessage.setVencryptVersion(message.majorVersion(), message.minorVersion() );
        sendMessage.writeBytes(m_connection);
    }

}

void ClientSocket::readVencryptSubAuthTypes()
{
    ServerVencryptAuthTypeResponseMessage message = ServerMessageBase::readMessage<ServerVencryptAuthTypeResponseMessage>(m_connection);
    if ( message.isVencryptVersionAccepted() )
    {
        qDebug() << "Subauth types:" << message.authTypes();
        if ( message.authTypes().contains( m_subAuthType ) )
        {
            ClientVencryptauthTypesMessage sendMessage;
            sendMessage.setSubauthType( m_subAuthType );
            sendMessage.writeBytes(m_connection);
        }
        else
            disconnectFromHost();
    }
    else
        disconnectFromHost();
}

void ClientSocket::readVencryptSubAuthResult()
{
    Uint8Reader result = Uint8Reader::readMessage<Uint8Reader>(m_connection);
    if ( result.value() == 0 )
    {
        disconnectFromHost();
    }
    else
    {
        qDebug("Start encryption.");
        m_connection->startClientEncryption();
    }

}

void ClientSocket::handleSslError( const QList<QSslError> & errors )
{
    qDebug() << "SSL Errors:" << errors;
    qDebug() << "Peer cert:" << m_connection->peerCertificate();
    qDebug() << "CA certs:" << m_connection->caCertificates();
}

void ClientSocket::handleEncryptionReady()
{
     qDebug() << "SSL Encrypted";
    if ( m_subAuthType == ClientGlobals::TLSNone )
        sendClientInit();
    else if ( m_subAuthType == ClientGlobals::TLSVNC )
    {
    }
    else
    {
    }
}

const QImage *ClientSocket::aquireImage() const
{
    return &m_framebuffer;
}

void ClientSocket::onDisconnected()
{

}

void ClientSocket::onHandshaking()
{

}

void ClientSocket::onConnected()
{

}

void ClientSocket::onSocketConnected()
{
    qDebug() << "Socket connected...";
}

void ClientSocket::onSocketError()
{
    qDebug() << m_connection->errorString();
}

void ClientSocket::onHandshakeReadVersion()
{
    readVersion();
}

void ClientSocket::onHandshakeSendVersionResponse()
{
    sendClientVersion();
}

void ClientSocket::onHandshakeVersionError()
{

}

void ClientSocket::onHandshakeReadSecurityTypes()
{
    readSecurityTypes();
    if ( m_canAuthenticate )
        emit securityTypeSupported();
    else
        emit securityTypeNotSupported();
}

void ClientSocket::onHandshakeSendSecurityType()
{
    ClientSecurityTypeMessage message;
    message.setSecurityType( m_authType );
    qDebug() << "Send security type: " << message;
    message.writeBytes(m_connection);

    if( m_authType == ClientGlobals::None )
    {
        if( m_maxClientVersion != ClientGlobals::RFB_38 )
            emit skipResponse();
        else
            emit useNoAuth();
    }
    else if ( m_authType == ClientGlobals::VNC )
    {
        emit useVNCAuth();
    }
    else if ( m_authType == ClientGlobals::VenCrypt )
    {
        emit useVencryptAuth();
    }
    else
        emit securityTypeNotSupported();
}

void ClientSocket::onHandshakeReadSecurityAuthResult()
{
    qDebug() << "onHandshakeReadSecurityAuthResult";
    readAuthResult();
    if ( m_authenticated )
        emit authSuccessful();
    else
        emit authError();
}

void ClientSocket::onHandshakeSecurityError()
{
    qDebug() << "Security error...";
}

void ClientSocket::onHandshakeSendVNCAuth()
{
    ServerAuthenticationChallengeMessage challenge = ServerMessageBase::readMessage<ServerAuthenticationChallengeMessage>( m_connection );
    ClientAuthenticationResponseMessage response;
    response.setChallenge( challenge.challenge());
    response.setPassword( m_password );
    response.writeBytes( m_connection );
    qDebug() << "onHandshakeSendVNCAuth" << response;
}

void ClientSocket::onHandshakeSendClientInit()
{
    qDebug() << "Send client init...";
    sendClientInit();
}

void ClientSocket::onHandshakeReadFramebufferSetup()
{
    qDebug() << "State ReadInitFramebuffer";
    readInitFramebuffer();
    // if there is an error emit framebufferError();
    sendPixelFormat();
    sendEncodings();
    QRect rect = QRect(QPoint(0,0), m_framebuffer.size());
    ClientRequestFramebufferUpdateMessage message;
    message.setUpdateRect(rect);
    message.writeBytes(m_connection);
    emit framebufferResize( m_framebuffer.size() );
    emit connected();

}

void ClientSocket::onHandshakeFramebufferError()
{

}

void ClientSocket::onHandleFramebufferUpdateState()
{
    //qDebug("framebuffer update...");
}

void ClientSocket::onSessionDone()
{
    qDebug() << "Session done...";
    emit disconnected();
}

void ClientSocket::onReadFrameBufferMessageTypeState()
{
    Uint8Reader messageType = Uint8Reader::readMessage<Uint8Reader>(m_connection);
    switch ( messageType.value() )
    {
      case 0:
        emit gotRfbUpdateMessage();
        break;
      case 1:
      case 2:
      case 3:
        qDebug() << "Other message" << messageType.value();
        emit finishedReadingMessage();
        // Nothing yet...
        break;
      case 255: // QEmu messages
          emit gotQemuUpdateMessage();
          break;
      default:
        qDebug() << "Unkown message:" << messageType.value();
        emit finishedReadingMessage();
      break;
    };
}

void ClientSocket::onReadQemuUpdateMessageState()
{
    Uint8Reader qemuMessage = Uint8Reader::readMessage<Uint8Reader>(m_connection);
    switch( qemuMessage.value() )
    {
      case 1: // Audio capture notify
      {
          emit gotAudioMessage();
      }
      break;
      default:
          qDebug("Unkonwn QEmu message %d", qemuMessage.value());
          emit finishedReadingMessage();
      break;
    };
}

void ClientSocket::onReadRfbUpdateMessageState()
{
    ConnectionProfiler::instance()->addFBUpdate();
    m_currentFbMessage = ServerMessageBase::readMessage<ServerFramebufferUpdateMessage>( m_connection );
    m_currentFbMessage.setPixelFormat( m_pixelFormat );
    m_currentFbMessage.setCompressionContext( m_unzip );
    m_currentFbMessage.setFramebufferSize( m_framebuffer.size() );
    m_currentFbMessage.setFramebufferDepth( m_framebuffer.depth() );
    if ( m_currentFbMessage.hasData() )
    {
        emit parseNextRfbUpdateRect();
    }
    else
    {        
        emit finishedReadingMessage();
    }
}

void ClientSocket::onReadRfbUpdateParseNextMessageState()
{

    ServerFramebufferUpdateMessage::FrameBufferUpdates update = m_currentFbMessage.readNextRectUpdate( m_connection );
    m_currentUpdateRect = update.rect;
    //qDebug() << "got rfb message" << update.encoding <<  update.rect << m_currentUpdateRect;

    switch( update.encoding )
    {
    case ClientGlobals::Raw:
        emit gotRectRawMessage();
        break;
    case ClientGlobals::RRE:
        emit gotRectRREMessage();
        break;
    case ClientGlobals::CoRRE:
        emit gotRectCoRREMessage();
        break;
    case ClientGlobals::CopyRect:
        emit gotRectCopyRectMessage();
        break;
    case ClientGlobals::Hextile:
        emit gotRectHextileMessage();
        break;
    case ClientGlobals::WMVi:
        emit gotWMViMessage();
        break;
    case ClientGlobals::PointerTypeChange:
        emit gotPointerTypeChangeMessage();
        break;
    case ClientGlobals::Audio:
        emit gotAudioMessage();
        break;
    case ClientGlobals::NativeKeys:
        emit gotNativeKeysMessage();
        break;
    case ClientGlobals::DesktopResize:
        emit gotDesktopResizeMessage();
        break;
    case ClientGlobals::Zlib:
        emit gotRectZlibMessage();
        break;
    case ClientGlobals::EncodingDone:
        emit framebufferUpdate();
        emit finishedReadingMessage();
        break;
    default:
        qWarning("Unknown encoding %d Rect: %d,%d - %dx%d",
                 update.encoding, update.rect.x(),
                 update.rect.y(), update.rect.width(),
                 update.rect.height() );
        emit finishedReadingMessage();
        break;
    }
}

void ClientSocket::onReadRfbUpdateRectDesktopResizeState()
{
    updateFrambufferFromPixelFormat( m_currentUpdateRect.size() );
    emit framebufferResize( m_currentUpdateRect.size() );
}

void ClientSocket::onReadRfbUpdateRectWMViState()
{
    m_pixelFormat = PixelFormat::readMessage<PixelFormat>(m_connection);
    updateFrambufferFromPixelFormat( m_framebuffer.size() );
    emit framebufferResize( m_currentUpdateRect.size() );
}

void ClientSocket::onReadRfbUpdateRectPointerTypeChangeState()
{
    emit enableRelativePointer();
}

void ClientSocket::onReadRfbUpdateRectAudioState()
{
    m_audioSupported = true;
    emit enableAudio();
}

void ClientSocket::onReadRfbUpdateRectNativeKeysState()
{
    emit enableNativeKeys();
}

void ClientSocket::onReadRfbUpdateRectRawState()
{
    qDebug() << "Raw" << m_currentUpdateRect;
    m_rawRectReadState->configure( m_framebuffer.rect(), m_pixelFormat.bpp(), m_fbBits, m_currentUpdateRect );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadRfbUpdateRectRREState()
{
    m_currentFbMessage.readRreRect( m_connection, m_fbBits, m_currentUpdateRect   );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadRfbUpdateRectCoRREState()
{
    m_currentFbMessage.readCoRreRect( m_connection, m_fbBits, m_currentUpdateRect   );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadRfbUpdateRectCopyRectState()
{
    m_currentFbMessage.readCopyRect( m_connection, m_fbBits, m_currentUpdateRect  );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadRfbUpdateRectHextileState()
{
    m_currentFbMessage.readHextileRect( m_connection, m_fbBits, m_currentUpdateRect  );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadRfbUpdateRectZlibState()
{
    m_currentFbMessage.readZlibRect( m_connection, m_fbBits, m_currentUpdateRect  );
    emit sendFramebufferChange( m_currentUpdateRect, m_framebuffer.copy(m_currentUpdateRect) );
}

void ClientSocket::onReadQemuUpdateAudioMessageState()
{
    ServerAudioCaptureNotifyMessage audioMessage = ServerAudioCaptureNotifyMessage::readMessage<ServerAudioCaptureNotifyMessage>( m_connection );
    if ( m_audioDevice )
    {
        if( audioMessage.audioState() == ServerAudioCaptureNotifyMessage::Connected && m_audioDevice->state() != QAudio::ActiveState )
        {
            m_audioBuffer = m_audioDevice->start();
        }
        else if( audioMessage.audioState() == ServerAudioCaptureNotifyMessage::Disconnected && m_audioDevice->state() == QAudio::ActiveState )
        {
            m_audioDevice->stop();
            m_audioBuffer = 0;
        }

        if ( !audioMessage.audioPacket().isEmpty() )
        {
            ConnectionProfiler::instance()->addAudioBytes( audioMessage.audioPacket().size() );
            m_audioBuffer->write( audioMessage.audioPacket() );
        }
    }
    emit finishedReadingMessage();
}

void ClientSocket::onHandleFramebufferUpdateDoneState()
{
    if ( m_connection->bytesAvailable() > 0 )
        emit moreBytesAvailable();        
}

void ClientSocket::onHandshakeReadVencryptVersion()
{
    qDebug() << "ClientSocket::onHandshakeReadVencryptVersion();";
    quint8 minor = 0;
    // read version and respond back with 0.2
    ServerVencryptVersionMessage version = ServerMessageBase::readMessage<ServerVencryptVersionMessage>(m_connection);
    if ( version.majorVersion() == 0 && version.minorVersion() == 2 )
    {
        minor = 2;
    }

    ClientVencryptVersionMessage result;
    result.setVencryptVersion(0,minor);
    result.writeBytes(m_connection);
}

void ClientSocket::onHandshakeReadVencryptSubauth()
{
    qDebug() << "ClientSocket::onHandshakeReadVencryptSubauth();" << m_subAuthType;
    ClientVencryptauthTypesMessage response;
    ServerVencryptAuthTypeResponseMessage authtypes = ServerMessageBase::readMessage<ServerVencryptAuthTypeResponseMessage>(m_connection);

    if( authtypes.isVencryptVersionAccepted() &&
        authtypes.authTypes().contains(m_subAuthType))
    {
        response.setSubauthType( m_subAuthType );
        response.writeBytes(m_connection);

        switch( m_subAuthType )
        {
        case ClientGlobals::Plain:
            emit usePlainAuth();
            break;
        case ClientGlobals::TLSNone:
            emit useTLS();
            break;
        case ClientGlobals::TLSVNC:
            emit useTLSVNC();
            break;
        case ClientGlobals::TLSPlain:
            emit useTLSPlain();
            break;
        case ClientGlobals::X509None:
            emit useX509();
            break;
        case ClientGlobals::X509VNC:
            emit useX509VNC();
            break;
        case ClientGlobals::X509Plain:
            emit useX509Plain();
            break;
        default:
            emit authError();
            break;
        }
    }
    else
    {
        response.setSubauthType( ClientGlobals::SubTypeInvalid );
        response.writeBytes(m_connection);
        emit authError();
    }
}

void ClientSocket::onHandshakeDoVencryptAuth()
{
    qDebug() << "ClientSocket::onHandshakeDoVencryptAuth()";
}

void ClientSocket::onHandshakeSendVencryptPlainAuth()
{
    qDebug() << "ClientSocket::onHandshakeSendVencryptPlainAuth()";
    ClientVeNCryptPlainAuthenticationMessage auth;
    auth.setUserInfo( m_username, m_password );
    auth.writeBytes(m_connection);
}

void ClientSocket::setupSSL( QSslSocket::PeerVerifyMode mode )
{
    Uint8Reader result = Uint8Reader::readMessage<Uint8Reader>(m_connection);
    if ( result.value() == 1 )
    {
        QSslError error(QSslError::SelfSignedCertificate);
        m_connection->ignoreSslErrors( QList<QSslError>() << error );
        m_connection->addCaCertificates(m_caCertPath);
        m_connection->setProtocol( QSsl::TlsV1 );
        m_connection->setPeerVerifyMode( mode );
        m_connection->startClientEncryption();
    }
    else
        emit authError();
}

void ClientSocket::onHandshakeStartTLS()
{
    qDebug() << "ClientSocket::onHandshakeStartTLS()";
    setupSSL(QSslSocket::QueryPeer);
}

void ClientSocket::onHandshakeStartTLSVNC()
{
    qDebug() << "ClientSocket::onHandshakeStartTLSVNC();";
    setupSSL(QSslSocket::QueryPeer);
}

void ClientSocket::onHandshakeStartTLSPlain()
{
    qDebug() << "ClientSocket::onHandshakeStartTLSPlain();";
    setupSSL(QSslSocket::QueryPeer);
}

void ClientSocket::onHandshakeStartX509()
{
    qDebug() << "ClientSocket::onHandshakeStartX509();";
    setupSSL(QSslSocket::VerifyPeer);
}

void ClientSocket::onHandshakeStartX509VNC()
{
    qDebug() << "ClientSocket::onHandshakeStartX509VNC();";
    setupSSL(QSslSocket::QueryPeer);
}

void ClientSocket::onHandshakeStartX509Plain()
{
    qDebug() << "ClientSocket::onHandshakeStartX509Plain();";
    setupSSL(QSslSocket::VerifyPeer);
}

void ClientSocket::onHandshakeReadVencryptAuthResult()
{
    qDebug() << "ClientSocket::onHandshakeReadVencryptAuthResult()";
    ServerVencryptAuthResponse response = ServerMessageBase::readMessage<ServerVencryptAuthResponse>(m_connection);
    if( response.success() )
        emit authSuccessful();
    else
        emit authError();
}