/****************************************************************************

**

** Copyright (C) 2010 Nokia Corporation and/or its subsidiary(-ies).

** All rights reserved.

** Contact: Nokia Corporation (qt-info@nokia.com)

**

** This file is part of the examples of the Qt Toolkit.

**

**

**

** $QT_END_LICENSE$

**

****************************************************************************/

#include "connectionmanager.h"

#include "filemanager.h"

#include "metainfo.h"

#include "torrentclient.h"

#include "torrentserver.h"

#include "trackerclient.h"

#include "peerwireclient.h"

#include "ratecontroller.h"

#include <QtCore>

#include <QNetworkInterface>

// These constants could also be configurable by the user.

static const int ServerMinPort = 6881;

static const int ServerMaxPort = /* 6889 */ 7000;

static const int BlockSize = 16384;

static const int MaxBlocksInProgress = 5;

static const int MaxBlocksInMultiMode = 2;

static const int MaxConnectionPerPeer = 1;

static const int RateControlWindowLength = 10;

static const int RateControlTimerDelay = 1000;

static const int MinimumTimeBeforeRevisit = 30;

static const int MaxUploads = 4;

static const int UploadScheduleInterval = 10000;

static const int EndGamePieces = 5;

class TorrentPiece {

public:

    int index;

    int length;

    QBitArray completedBlocks;

    QBitArray requestedBlocks;

    bool inProgress;

};

class TorrentClientPrivate

{

public:

    TorrentClientPrivate(TorrentClient *qq);

    // State / error

    void setError(TorrentClient::Error error);

    void setState(TorrentClient::State state);

    TorrentClient::Error error;

    TorrentClient::State state;

    QString errorString;

    QString stateString;

    // Where to save data

    QString destinationFolder;

    MetaInfo metaInfo;

    // Announce tracker and file manager

    QByteArray peerId;

    QByteArray infoHash;

    TrackerClient trackerClient;

    FileManager fileManager;

    // Connections

    QList<PeerWireClient *> connections;

    QList<TorrentPeer *> peers;

    bool schedulerCalled;

    void callScheduler();

    bool connectingToClients;

    void callPeerConnector();

    int uploadScheduleTimer;

    // Pieces

    QMap<int, PeerWireClient *> readIds;

    QMultiMap<PeerWireClient *, TorrentPiece *> payloads;

    QMap<int, TorrentPiece *> pendingPieces;

    QBitArray completedPieces;

    QBitArray incompletePieces;

    int pieceCount;

    // Progress

    int lastProgressValue;

    qint64 downloadedBytes;

    qint64 uploadedBytes;

    int downloadRate[RateControlWindowLength];

    int uploadRate[RateControlWindowLength];

    int transferRateTimer;

    TorrentClient *q;

};

TorrentClientPrivate::TorrentClientPrivate(TorrentClient *qq)

    : trackerClient(qq), q(qq)

{

    error = TorrentClient::UnknownError;

    state = TorrentClient::Idle;

    errorString = QT_TRANSLATE_NOOP(TorrentClient, "Unknown error");

    stateString = QT_TRANSLATE_NOOP(TorrentClient, "Idle");

    schedulerCalled = false;

    connectingToClients = false;

    uploadScheduleTimer = 0;

    lastProgressValue = -1;

    pieceCount = 0;

    downloadedBytes = 0;

    uploadedBytes = 0;

    memset(downloadRate, 0, sizeof(downloadRate));

    memset(uploadRate, 0, sizeof(uploadRate));

    transferRateTimer = 0;

}

void TorrentClientPrivate::setError(TorrentClient::Error errorCode)

{

    this->error = errorCode;

    switch (error) {

    case TorrentClient::UnknownError:

        errorString = QT_TRANSLATE_NOOP(TorrentClient, "Unknown error");

        break;

    case TorrentClient::TorrentParseError:

        errorString = QT_TRANSLATE_NOOP(TorrentClient, "Invalid torrent data");

        break;

    case TorrentClient::InvalidTrackerError:

        errorString = QT_TRANSLATE_NOOP(TorrentClient, "Unable to connect to tracker");

        break;

    case TorrentClient::FileError:

        errorString = QT_TRANSLATE_NOOP(TorrentClient, "File error");

        break;

    case TorrentClient::ServerError:

        errorString = QT_TRANSLATE_NOOP(TorrentClient, "Unable to initialize server");

        break;

    }

    emit q->error(errorCode);

}

void TorrentClientPrivate::setState(TorrentClient::State state)

{

    this->state = state;

    switch (state) {

    case TorrentClient::Idle:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Idle");

        break;

    case TorrentClient::Paused:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Paused");

        break;

    case TorrentClient::Stopping:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Stopping");

        break;

    case TorrentClient::Preparing:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Preparing");

        break;

    case TorrentClient::Searching:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Searching");

        break;

    case TorrentClient::Connecting:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Connecting");

        break;

    case TorrentClient::WarmingUp:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Warming up");

        break;

    case TorrentClient::Downloading:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Downloading");

        break;

    case TorrentClient::Endgame:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Finishing");

        break;

    case TorrentClient::Seeding:

        stateString = QT_TRANSLATE_NOOP(TorrentClient, "Seeding");

        break;

    }

    emit q->stateChanged(state);

}

void TorrentClientPrivate::callScheduler()

{

    if (!schedulerCalled) {

        schedulerCalled = true;

        QMetaObject::invokeMethod(q, "scheduleDownloads", Qt::QueuedConnection);

    }

}

void TorrentClientPrivate::callPeerConnector()

{

    if (!connectingToClients) {

        connectingToClients = true;

        QTimer::singleShot(10000, q, SLOT(connectToPeers()));

    }

}

TorrentClient::TorrentClient(QObject *parent)

    : QObject(parent), d(new TorrentClientPrivate(this))

{

    // Connect the file manager

    connect(&d->fileManager, SIGNAL(dataRead(int,int,int,QByteArray)),

            this, SLOT(sendToPeer(int,int,int,QByteArray)));

    connect(&d->fileManager, SIGNAL(verificationProgress(int)),

            this, SLOT(updateProgress(int)));

    connect(&d->fileManager, SIGNAL(verificationDone()),

            this, SLOT(fullVerificationDone()));

    connect(&d->fileManager, SIGNAL(pieceVerified(int,bool)),

            this, SLOT(pieceVerified(int,bool)));

    connect(&d->fileManager, SIGNAL(error()),

            this, SLOT(handleFileError()));

    // Connect the tracker client

    connect(&d->trackerClient, SIGNAL(peerListUpdated(QList<TorrentPeer>)),

            this, SLOT(addToPeerList(QList<TorrentPeer>)));

    connect(&d->trackerClient, SIGNAL(stopped()),

            this, SIGNAL(stopped()));

}

TorrentClient::~TorrentClient()

{

    qDeleteAll(d->peers);

    qDeleteAll(d->pendingPieces);

    delete d;

}

bool TorrentClient::setTorrent(const QString &fileName)

{

    QFile file(fileName);

    if (!file.open(QIODevice::ReadOnly) || !setTorrent(file.readAll())) {

        d->setError(TorrentParseError);

        return false;

    }

    return true;

}

bool TorrentClient::setTorrent(const QByteArray &torrentData)

{

    if (!d->metaInfo.parse(torrentData)) {

        d->setError(TorrentParseError);

        return false;

    }

    // Calculate SHA1 hash of the "info" section in the torrent

    QByteArray infoValue = d->metaInfo.infoValue();

    d->infoHash = QCryptographicHash::hash(infoValue, QCryptographicHash::Sha1);

    return true;

}

MetaInfo TorrentClient::metaInfo() const

{

    return d->metaInfo;

}

void TorrentClient::setDestinationFolder(const QString &directory)

{

    d->destinationFolder = directory;

}

QString TorrentClient::destinationFolder() const

{

    return d->destinationFolder;

}

void TorrentClient::setDumpedState(const QByteArray &dumpedState)

{

    // Recover partially completed pieces

    QDataStream stream(dumpedState);

    quint16 version = 0;

    stream >> version;

    if (version != 2)

        return;

    stream >> d->completedPieces;

    while (!stream.atEnd()) {

        int index;

        int length;

        QBitArray completed;

        stream >> index >> length >> completed;

        if (stream.status() != QDataStream::Ok) {

            d->completedPieces.clear();

            break;

        }

        TorrentPiece *piece = new TorrentPiece;

        piece->index = index;

        piece->length = length;

        piece->completedBlocks = completed;

        piece->requestedBlocks.resize(completed.size());

        piece->inProgress = false;

        d->pendingPieces[index] = piece;

    }

}

QByteArray TorrentClient::dumpedState() const

{

    QByteArray partials;

    QDataStream stream(&partials, QIODevice::WriteOnly);

    stream << quint16(2);

    stream << d->completedPieces;

    // Save the state of all partially downloaded pieces into a format

    // suitable for storing in settings.

    QMap<int, TorrentPiece *>::ConstIterator it = d->pendingPieces.constBegin();

    while (it != d->pendingPieces.constEnd()) {

        TorrentPiece *piece = it.value();

        if (blocksLeftForPiece(piece) > 0 && blocksLeftForPiece(piece) < piece->completedBlocks.size()) {

            stream << piece->index;

            stream << piece->length;

            stream << piece->completedBlocks;

        }

        ++it;

    }

    return partials;

}

qint64 TorrentClient::progress() const

{

    return d->lastProgressValue;

}

void TorrentClient::setDownloadedBytes(qint64 bytes)

{

    d->downloadedBytes = bytes;

}

qint64 TorrentClient::downloadedBytes() const

{

    return d->downloadedBytes;

}

void TorrentClient::setUploadedBytes(qint64 bytes)

{

    d->uploadedBytes = bytes;

}

qint64 TorrentClient::uploadedBytes() const

{

    return d->uploadedBytes;

}

int TorrentClient::connectedPeerCount() const

{

    int tmp = 0;

    foreach (PeerWireClient *client, d->connections) {

        if (client->state() == QAbstractSocket::ConnectedState)

            ++tmp;

    }

    return tmp;

}

int TorrentClient::seedCount() const

{

    int tmp = 0;

    foreach (PeerWireClient *client, d->connections) {

        if (client->availablePieces().count(true) == d->pieceCount)

            ++tmp;

    }

    return tmp;

}

TorrentClient::State TorrentClient::state() const

{

    return d->state;

}

QString TorrentClient::stateString() const

{

    return d->stateString;

}

TorrentClient::Error TorrentClient::error() const

{

    return d->error;

}

QString TorrentClient::errorString() const

{

    return d->errorString;

}

QByteArray TorrentClient::peerId() const

{

    return d->peerId;

}

QByteArray TorrentClient::infoHash() const

{

    return d->infoHash;

}

void TorrentClient::start()

{

    if (d->state != Idle)

        return;

    TorrentServer::instance()->addClient(this);

    // Initialize the file manager

    d->setState(Preparing);

    d->fileManager.setMetaInfo(d->metaInfo);

    d->fileManager.setDestinationFolder(d->destinationFolder);

    d->fileManager.setCompletedPieces(d->completedPieces);

    d->fileManager.start(QThread::LowestPriority);

    d->fileManager.startDataVerification();

}

void TorrentClient::stop()

{

    if (d->state == Stopping)

        return;

    TorrentServer::instance()->removeClient(this);

    // Update the state

    State oldState = d->state;

    d->setState(Stopping);

    // Stop the timer

    if (d->transferRateTimer) {

        killTimer(d->transferRateTimer);

        d->transferRateTimer = 0;

    }

    // Abort all existing connections

    foreach (PeerWireClient *client, d->connections) {

        RateController::instance()->removeSocket(client);

        ConnectionManager::instance()->removeConnection(client);

        client->abort();

    }

    d->connections.clear();

    // Perhaps stop the tracker

    if (oldState > Preparing) {

        d->trackerClient.stop();

    } else {

        d->setState(Idle);

        emit stopped();

    }

}

void TorrentClient::setPaused(bool paused)

{

    if (paused) {

        // Abort all connections, and set the max number of

        // connections to 0. Keep the list of peers, so we can quickly

        // resume later.

        d->setState(Paused);

        foreach (PeerWireClient *client, d->connections)

            client->abort();

        d->connections.clear();

        TorrentServer::instance()->removeClient(this);

    } else {

        // Restore the max number of connections, and start the peer

        // connector. We should also quickly start receiving incoming

        // connections.

        d->setState(d->completedPieces.count(true) == d->fileManager.pieceCount()

                    ? Seeding : Searching);

        connectToPeers();

        TorrentServer::instance()->addClient(this);

    }

}

void TorrentClient::timerEvent(QTimerEvent *event)

{

    if (event->timerId() == d->uploadScheduleTimer) {

        // Update the state of who's choked and who's not

        scheduleUploads();

        return;

    }

    if (event->timerId() != d->transferRateTimer) {

        QObject::timerEvent(event);

        return;

    }

    // Calculate average upload/download rate

    qint64 uploadBytesPerSecond = 0;

    qint64 downloadBytesPerSecond = 0;

    for (int i = 0; i < RateControlWindowLength; ++i) {

        uploadBytesPerSecond += d->uploadRate[i];

        downloadBytesPerSecond += d->downloadRate[i];

    }

    uploadBytesPerSecond /= qint64(RateControlWindowLength);

    downloadBytesPerSecond /= qint64(RateControlWindowLength);

    for (int i = RateControlWindowLength - 2; i >= 0; --i) {

        d->uploadRate[i + 1] = d->uploadRate[i];

        d->downloadRate[i + 1] = d->downloadRate[i];

    }

    d->uploadRate[0] = 0;

    d->downloadRate[0] = 0;

    emit uploadRateUpdated(int(uploadBytesPerSecond));

    emit downloadRateUpdated(int(downloadBytesPerSecond));

    // Stop the timer if there is no activity.