/***************************************************************************
 *   Copyright (C) 2005 by Joris Guisson                                   *
 *   joris.guisson@gmail.com                                               *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.             *
 ***************************************************************************/
#include "peer.h"

#include <cmath>
#include <version.h>
#include <util/log.h>
#include <util/functions.h>
#include <net/address.h>
#include <mse/encryptedpacketsocket.h>
#include <diskio/chunk.h>
#include <download/piece.h>
#include <download/request.h>
#include <bcodec/bdecoder.h>
#include <bcodec/bencoder.h>
#include <bcodec/bnode.h>
#include <torrent/server.h>
#include <torrent/torrent.h>
#include "packetreader.h"
#include "peerdownloader.h"
#include "peeruploader.h"
#include "utpex.h"
#include "peermanager.h"
#include <net/reverseresolver.h>
#include "utmetadata.h"

using namespace net;

namespace bt
{

static const int MAX_METADATA_SIZE = 4 * 1024 * 1024; // maximum allowed metadata_size (up to 4 MiB)
static Uint32 peer_id_counter = 1;
bool Peer::resolve_hostname = true;


Peer::Peer(mse::EncryptedPacketSocket::Ptr sock,
           const PeerID & peer_id,
           Uint32 num_chunks,
           Uint32 chunk_size,
           Uint32 support,
           bool local,
           ConnectionLimit::Token::Ptr token,
           PeerManager* pman)
    : PeerInterface(peer_id, num_chunks),
      sock(sock),
      token(token),
      pman(pman)
{
    id = peer_id_counter;
    peer_id_counter++;
    ut_pex_id = 0;

    Uint32 max_packet_len = 9 + MAX_PIECE_LEN; // size of piece packet
    Uint32 bitfield_length = num_chunks / 8 + 1 + (num_chunks % 8 == 0 ? 0 : 1);
    if (bitfield_length > max_packet_len) // bitfield can be longer
        max_packet_len = bitfield_length;

    // to be future proof use 10 times the max packet length,
    // in case some extension message comes along which is larger
    preader = new PacketReader(10 * max_packet_len);
    downloader = new PeerDownloader(this, chunk_size);
    uploader = new PeerUploader(this);

    stalled_timer.update();


    connect_time = QTime::currentTime();
    stats.client = peer_id.identifyClient();
    stats.ip_address = getIPAddresss();
    stats.dht_support = support & DHT_SUPPORT;
    stats.fast_extensions = support & FAST_EXT_SUPPORT;
    stats.extension_protocol = support & EXT_PROT_SUPPORT;
    stats.encrypted = sock->encrypted();
    stats.local = local;
    stats.transport_protocol = sock->socketDevice()->transportProtocol();

    if (stats.ip_address == QStringLiteral("0.0.0.0")) {
        Out(SYS_CON | LOG_DEBUG) << "No more 0.0.0.0" << endl;
        kill();
    } else {
        sock->startMonitoring(preader);
    }

    pex_allowed = stats.extension_protocol;
    extensions.setAutoDelete(true);

    if (resolve_hostname) {
        net::ReverseResolver* res = new net::ReverseResolver();
        connect(res, &net::ReverseResolver::resolved, this, &Peer::resolved, Qt::QueuedConnection);
        res->resolveAsync(sock->getRemoteAddress());
    }
}


Peer::~Peer()
{
    // Seeing that we are going to delete the preader object,
    // call stopMonitoring, in some situations it is possible that the socket
    // is only deleted later because the authenticate object still has a reference to it
    sock->stopMonitoring();
    delete uploader;
    delete downloader;
    delete preader;
}

void Peer::closeConnection()
{
    sock->close();
}

void Peer::kill()
{
    sock->close();
    killed = true;
    token.clear();
}

void Peer::handleChoke(Uint32 len)
{
    if (len != 1) {
        kill();
        return;
    }

    if (!stats.choked)
        stats.time_choked = CurrentTime();
    stats.choked = true;
    downloader->choked();
}

void Peer::handleUnchoke(Uint32 len)
{
    if (len != 1) {
        Out(SYS_CON | LOG_DEBUG) << "len err UNCHOKE" << endl;
        kill();
        return;
    }

    if (stats.choked) {
        stats.time_unchoked = CurrentTime();
        bytes_downloaded_since_unchoke = 0;
    }

    stats.choked = false;
}

void Peer::handleInterested(Uint32 len)
{
    if (len != 1) {
        Out(SYS_CON | LOG_DEBUG) << "len err INTERESTED" << endl;
        kill();
        return;
    }

    if (!stats.interested) {
        stats.interested = true;
        pman->rerunChoker();
    }

}

void Peer::handleNotInterested(Uint32 len)
{
    if (len != 1) {
        kill();
        return;
    }

    if (stats.interested) {
        stats.interested = false;
        pman->rerunChoker();
    }
}

void Peer::handleHave(const bt::Uint8* packet, Uint32 len)
{
    if (len != 5) {
        kill();
    } else {
        Uint32 ch = ReadUint32(packet, 1);
        if (ch < pieces.getNumBits()) {
            pman->have(this, ch);
            pieces.set(ch, true);
        } else if (pman->getTorrent().isLoaded()) {
            Out(SYS_CON | LOG_NOTICE) << "Received invalid have value, kicking peer" << endl;
            kill();
        }
    }
}

void Peer::handleHaveAll(Uint32 len)
{
    if (len != 1) {
        kill();
    } else {
        pieces.setAll(true);
        pman->bitSetReceived(this, pieces);
    }
}

void Peer::handleHaveNone(Uint32 len)
{
    if (len != 1) {
        kill();
    } else {
        pieces.setAll(false);
        pman->bitSetReceived(this, pieces);
    }
}

void Peer::handleBitField(const bt::Uint8* packet, Uint32 len)
{
    if (len != 1 + pieces.getNumBytes()) {
        if (pman->getTorrent().isLoaded())
            kill();
    } else {
        pieces = BitSet(packet + 1, pieces.getNumBits());
        pman->bitSetReceived(this, pieces);
    }
}

void Peer::handleRequest(const bt::Uint8* packet, Uint32 len)
{
    if (len != 13) {
        kill();
        return;
    }

    Request r(
        ReadUint32(packet, 1),
        ReadUint32(packet, 5),
        ReadUint32(packet, 9),
        downloader);

    if (stats.has_upload_slot)
        uploader->addRequest(r);
    else if (stats.fast_extensions)
        sendReject(r);
}

void Peer::handlePiece(const bt::Uint8* packet, Uint32 len)
{
    if (paused)
        return;

    if (len < 9) {
        kill();
        return;
    }

    snub_timer.update();

    stats.bytes_downloaded += (len - 9);
    bytes_downloaded_since_unchoke += (len - 9);
    Piece p(ReadUint32(packet, 1),
            ReadUint32(packet, 5),
            len - 9, downloader, packet + 9);
    downloader->piece(p);
    pman->pieceReceived(p);
    downloader->update();
}

void Peer::handleCancel(const bt::Uint8* packet, Uint32 len)
{
    if (len != 13) {
        kill();
    } else {
        Request r(ReadUint32(packet, 1),
                  ReadUint32(packet, 5),
                  ReadUint32(packet, 9),
                  downloader);
        uploader->removeRequest(r);
        sock->doNotSendPiece(r, stats.fast_extensions);
    }
}

void Peer::handleReject(const bt::Uint8* packet, Uint32 len)
{
    if (len != 13) {
        kill();
    } else {
        Request r(ReadUint32(packet, 1),
                  ReadUint32(packet, 5),
                  ReadUint32(packet, 9),
                  downloader);
        downloader->onRejected(r);
    }
}

void Peer::handlePort(const bt::Uint8* packet, Uint32 len)
{
    if (len != 3) {
        kill();
    } else {
        Uint16 port = ReadUint16(packet, 1);
        //  Out(SYS_CON|LOG_DEBUG) << "Got PORT packet : " << port << endl;
        pman->portPacketReceived(getIPAddresss(), port);
    }
}

void Peer::handlePacket(const bt::Uint8* packet, Uint32 size)
{
    if (killed || size == 0)
        return;

    switch (packet[0]) {
    case CHOKE:
        handleChoke(size);
        break;
    case UNCHOKE:
        handleUnchoke(size);
        break;
    case INTERESTED:
        handleInterested(size);
        break;
    case NOT_INTERESTED:
        handleNotInterested(size);
        break;
    case HAVE:
        handleHave(packet, size);
        break;
    case BITFIELD:
        handleBitField(packet, size);
        break;
    case REQUEST:
        handleRequest(packet, size);
        break;
    case PIECE:
        handlePiece(packet, size);
        break;
    case CANCEL:
        handleCancel(packet, size);
        break;
    case REJECT_REQUEST:
        handleReject(packet, size);
        break;
    case PORT:
        handlePort(packet, size);
        break;
    case HAVE_ALL:
        handleHaveAll(size);
        break;
    case HAVE_NONE:
        handleHaveNone(size);
        break;
    case SUGGEST_PIECE:
        // ignore suggestions for the moment
        break;
    case ALLOWED_FAST:
        // we no longer support this, so do nothing
        break;
    case EXTENDED:
        handleExtendedPacket(packet, size);
        break;
    }
}

void Peer::pause()
{
    if (paused)
        return;

    downloader->cancelAll();
    // choke the peer and tell it we are not interested
    choke();
    sendNotInterested();
    paused = true;
}

void Peer::unpause()
{
    paused = false;
}

void Peer::handleExtendedPacket(const Uint8* packet, Uint32 size)
{
    if (size <= 2)
        return;

    PeerProtocolExtension* ext = extensions.find(packet[1]);
    if (ext) {
        ext->handlePacket(packet, size);
    } else if (packet[1] == 0) {
        handleExtendedHandshake(packet, size);
    }
}

void Peer::handleExtendedHandshake(const Uint8* packet, Uint32 size)
{
    QByteArray tmp = QByteArray::fromRawData((const char*)packet, size);
    BNode* node = 0;
    try {
        BDecoder dec(tmp, false, 2);
        node = dec.decode();
        if (!node || node->getType() != BNode::DICT) {
            delete node;
            return;
        }

        BDictNode* dict = (BDictNode*)node;
        BDictNode* mdict = dict->getDict(QByteArrayLiteral("m"));
        if (!mdict) {
            delete node;
            return;
        }

        BValueNode* val = 0;

        if ((val = mdict->getValue(QByteArrayLiteral("ut_pex"))) && UTPex::isEnabled()) {
            // ut_pex packet
            ut_pex_id = val->data().toInt();
            if (ut_pex_id == 0) {
                extensions.erase(UT_PEX_ID);
            } else {
                PeerProtocolExtension* ext = extensions.find(UT_PEX_ID);
                if (ext)
                    ext->changeID(ut_pex_id);
                else if (pex_allowed)
                    extensions.insert(UT_PEX_ID, new UTPex(this, ut_pex_id));
            }
        }

        if ((val = mdict->getValue(QByteArrayLiteral("ut_metadata")))) {
            // meta data
            Uint32 ut_metadata_id = val->data().toInt();
            if (ut_metadata_id == 0) { // disabled by other side
                extensions.erase(UT_METADATA_ID);
            } else {
                PeerProtocolExtension* ext = extensions.find(UT_METADATA_ID);
                if (ext) {
                    ext->changeID(ut_metadata_id);
                } else {
                    int metadata_size = 0;
                    if (dict->getValue(QByteArrayLiteral("metadata_size")))
                        metadata_size = dict->getInt(QByteArrayLiteral("metadata_size"));

                    if (metadata_size > MAX_METADATA_SIZE) { // check for wrong metadata_size values
                        Out(SYS_GEN | LOG_NOTICE) << "Wrong metadata size: " << metadata_size << ". Killing peer... " << endl;
                        kill();
                        delete node;
                        return;
                    }

                    UTMetaData* md = new UTMetaData(pman->getTorrent(), ut_metadata_id, this);
                    md->setReportedMetadataSize(metadata_size);
                    extensions.insert(UT_METADATA_ID, md);
                }
            }
        }

        if ((val = dict->getValue(QByteArrayLiteral("reqq")))) {
            stats.max_request_queue = val->data().toInt();
        }

        if ((val = dict->getValue(QByteArrayLiteral("upload_only")))) {
            stats.partial_seed = val->data().toInt() == 1;
        }
    } catch (...) {
        // just ignore invalid packets
        Out(SYS_CON | LOG_DEBUG) << "Invalid extended packet" << endl;
    }
    delete node;
}

Uint32 Peer::sendData(const Uint8* data, Uint32 len)
{
    if (killed) return 0;

    Uint32 ret = sock->sendData(data, len);
    if (!sock->ok())
        kill();

    return ret;
}

Uint32 Peer::readData(Uint8* buf, Uint32 len)
{
    if (killed) return 0;

    Uint32 ret = sock->readData(buf, len);

    if (!sock->ok())
        kill();

    return ret;
}

Uint32 Peer::bytesAvailable() const
{
    return sock->bytesAvailable();
}

Uint32 Peer::getUploadRate() const
{
    if (sock)
        return sock->getUploadRate();
    else
        return 0;
}

Uint32 Peer::getDownloadRate() const
{
    if (sock)
        return sock->getDownloadRate();
    else
        return 0;
}

void Peer::update()
{
    if (killed)
        return;

    if (!sock->ok() || !preader->ok()) {
        Out(SYS_CON | LOG_DEBUG) << "Connection closed" << endl;
        kill();
        return;
    }

    sock->updateSpeeds(bt::CurrentTime());
    preader->update(*this);

    Uint32 data_bytes = sock->dataBytesUploaded();

    if (data_bytes > 0) {
        stats.bytes_uploaded += data_bytes;
        uploader->addUploadedBytes(data_bytes);
    }

    if (!paused) {
        PtrMap<Uint32, PeerProtocolExtension>::iterator i = extensions.begin();
        while (i != extensions.end()) {
            if (i->second->needsUpdate())
                i->second->update();
            ++i;
        }
    }

    // if no data is being sent or received, and there are pending requests
    // increment the connection stalled timer
    if (getUploadRate() > 100 || getDownloadRate() > 100 ||
        (uploader->getNumRequests() == 0 && sock->numPendingPieceUploads() == 0 && downloader->getNumRequests() == 0))
        stalled_timer.update();

    stats.download_rate = this->getDownloadRate();
    stats.upload_rate = this->getUploadRate();
    stats.perc_of_file = this->percentAvailable();
    stats.snubbed = this->isSnubbed();
    stats.num_up_requests = uploader->getNumRequests() + sock->numPendingPieceUploads();
    stats.num_down_requests = downloader->getNumRequests();
}

bool Peer::isStalled() const
{
    return stalled_timer.getElapsedSinceUpdate() >= 2 * 60 * 1000;
}

bool Peer::isSnubbed() const
{
    // 4 minutes
    return snub_timer.getElapsedSinceUpdate() >= 2 * 60 * 1000 && stats.num_down_requests > 0;
}

QString Peer::getIPAddresss() const
{
    if (sock)
        return sock->getRemoteIPAddress();
    else
        return QString();
}

Uint16 Peer::getPort() const
{
    if (!sock)
        return 0;
    else
        return sock->getRemotePort();
}

net::Address Peer::getAddress() const
{
    if (!sock)
        return net::Address();
    else
        return sock->getRemoteAddress();
}

Uint32 Peer::getTimeSinceLastPiece() const
{
    return snub_timer.getElapsedSinceUpdate();
}

float Peer::percentAvailable() const
{
    // calculation needs to use bytes, instead of chunks, because
    // the last chunk can have a different size
    const Torrent & tor = pman->getTorrent();
    Uint64 bytes = 0;
    if (pieces.get(tor.getNumChunks() - 1))
        bytes = tor.getChunkSize() * (pieces.numOnBits() - 1) + tor.getLastChunkSize();
    else
        bytes = tor.getChunkSize() * pieces.numOnBits();

    Uint64 tbytes = tor.getChunkSize() * (pieces.getNumBits() - 1) + tor.getLastChunkSize();
    return (float)bytes / (float)tbytes * 100.0;
}

void Peer::setACAScore(double s)
{
    stats.aca_score = s;
}

void Peer::choke()
{
    if (!stats.has_upload_slot)
        return;

    sendChoke();
    uploader->clearAllRequests();
}

void Peer::emitPortPacket()
{
    pman->portPacketReceived(sock->getRemoteIPAddress(), sock->getRemotePort());
}

void Peer::emitPex(const QByteArray & data)
{
    pman->pex(data);
}

void Peer::setPexEnabled(bool on)
{
    if (!stats.extension_protocol)
        return;

    PeerProtocolExtension* ext = extensions.find(UT_PEX_ID);
    if (ext && (!on || !UTPex::isEnabled())) {
        extensions.erase(UT_PEX_ID);
    } else if (!ext && on && ut_pex_id > 0 && UTPex::isEnabled()) {
        // if the other side has enabled it to, create a new UTPex object
        extensions.insert(UT_PEX_ID, new UTPex(this, ut_pex_id));
    }

    pex_allowed = on;
}

void Peer::sendExtProtHandshake(Uint16 port, Uint32 metadata_size, bool partial_seed)
{
    if (!stats.extension_protocol)
        return;

    QByteArray arr;
    BEncoder enc(new BEncoderBufferOutput(arr));
    enc.beginDict();
    enc.write(QByteArrayLiteral("m"));
    // supported messages
    enc.beginDict();
    enc.write(QByteArrayLiteral("ut_pex")); enc.write((Uint32)(pex_allowed ? UT_PEX_ID : 0));
    enc.write(QByteArrayLiteral("ut_metadata")); enc.write(UT_METADATA_ID);
    enc.end();
    if (port > 0) {
        enc.write(QByteArrayLiteral("p"));
        enc.write((Uint32)port);
    }
    enc.write(QByteArrayLiteral("reqq"));
    enc.write((bt::Uint32)250);
    if (metadata_size) {
        enc.write(QByteArrayLiteral("metadata_size"));
        enc.write(metadata_size);
    }

    enc.write(QByteArrayLiteral("upload_only"), partial_seed ? QByteArrayLiteral("1") : QByteArrayLiteral("0"));
    enc.write(QByteArrayLiteral("v")); enc.write(bt::GetVersionString().toLatin1());
    enc.end();
    sendExtProtMsg(0, arr);
}


void Peer::setGroupIDs(Uint32 up_gid, Uint32 down_gid)
{
    sock->setGroupID(up_gid, true);
    sock->setGroupID(down_gid, false);
}

void Peer::resolved(const QString & hinfo)
{
    stats.hostname = hinfo;
}

void Peer::setResolveHostnames(bool on)
{
    resolve_hostname = on;
}

void Peer::emitMetadataDownloaded(const QByteArray& data)
{
    Q_EMIT metadataDownloaded(data);
}

bool Peer::hasWantedChunks(const bt::BitSet& wanted_chunks) const
{
    BitSet bs = pieces;
    bs.andBitSet(wanted_chunks);
    return bs.numOnBits() > 0;
}

Uint32 Peer::averageDownloadSpeed() const
{
    if (stats.choked)
        return 0;

    TimeStamp now = CurrentTime();
    if (now >= getUnchokeTime())
        return 0;
    else
        return bytes_downloaded_since_unchoke / (getUnchokeTime() - now);
}


void Peer::sendChoke()
{
    if (!stats.has_upload_slot)
        return;

    sock->addPacket(Packet::Ptr(new Packet(CHOKE)));
    stats.has_upload_slot = false;
}

void Peer::sendUnchoke()
{
    if (stats.has_upload_slot)
        return;

    sock->addPacket(Packet::Ptr(new Packet(UNCHOKE)));
    stats.has_upload_slot = true;
}

void Peer::sendEvilUnchoke()
{
    sock->addPacket(Packet::Ptr(new Packet(UNCHOKE)));
    stats.has_upload_slot = false;
}

void Peer::sendInterested()
{
    if (stats.am_interested == true)
        return;

    sock->addPacket(Packet::Ptr(new Packet(INTERESTED)));
    stats.am_interested = true;
}

void Peer::sendNotInterested()
{
    if (stats.am_interested == false)
        return;

    sock->addPacket(Packet::Ptr(new Packet(NOT_INTERESTED)));
    stats.am_interested = false;
}

void Peer::sendRequest(const Request & r)
{
    sock->addPacket(Packet::Ptr(new Packet(r, bt::REQUEST)));
}

void Peer::sendCancel(const Request & r)
{
    sock->addPacket(Packet::Ptr(new Packet(r, bt::CANCEL)));
}

void Peer::sendReject(const Request & r)
{
    sock->addPacket(Packet::Ptr(new Packet(r, bt::REJECT_REQUEST)));
}

void Peer::sendHave(Uint32 index)
{
    sock->addPacket(Packet::Ptr(new Packet(index, bt::HAVE)));
}

void Peer::sendPort(Uint16 port)
{
    sock->addPacket(Packet::Ptr(new Packet(port)));
}

void Peer::sendBitSet(const BitSet & bs)
{
    sock->addPacket(Packet::Ptr(new Packet(bs)));
}

void Peer::sendHaveAll()
{
    sock->addPacket(Packet::Ptr(new Packet(bt::HAVE_ALL)));
}

void Peer::sendHaveNone()
{
    sock->addPacket(Packet::Ptr(new Packet(bt::HAVE_NONE)));
}

void Peer::sendSuggestPiece(Uint32 index)
{
    sock->addPacket(Packet::Ptr(new Packet(index, bt::SUGGEST_PIECE)));
}

void Peer::sendAllowedFast(Uint32 index)
{
    sock->addPacket(Packet::Ptr(new Packet(index, bt::ALLOWED_FAST)));
}

bool Peer::sendChunk(Uint32 index, Uint32 begin, Uint32 len, Chunk * ch)
{
    //      Out() << "sendChunk " << index << " " << begin << " " << len << endl;

#ifndef NDEBUG
    Q_ASSERT(ch);
    if (!ch) {
        Out(SYS_CON | LOG_NOTICE) << "Warning : attempted to upload an invalid chunk" << endl;
        return false;
    }
#endif
    if (begin >= ch->getSize() || begin + len > ch->getSize()) {
        Out(SYS_CON | LOG_NOTICE) << "Warning : Illegal piece request" << endl;
        Out(SYS_CON | LOG_NOTICE) << "\tChunk : index " << index << " size = " << ch->getSize() << endl;
        Out(SYS_CON | LOG_NOTICE) << "\tPiece : begin = " << begin << " len = " << len << endl;
        return false;
    }
    /*      Out(SYS_CON|LOG_DEBUG) << QString("Uploading %1 %2 %3 %4 %5")
     *          .arg(index).arg(begin).arg(len).arg((quint64)ch,0,16).arg((quint64)ch->getData(),0,16)
     *          << endl;;
     */
    sock->addPacket(Packet::Ptr(new Packet(index, begin, len, ch)));
    return true;
}

void Peer::sendExtProtMsg(Uint8 id, const QByteArray & data)
{
    sock->addPacket(Packet::Ptr(new Packet(id, data)));
}

void Peer::clearPendingPieceUploads()
{
    sock->clearPieces(stats.fast_extensions);
}

void Peer::chunkAllowed(Uint32 chunk)
{
    sendHave(chunk);
}

}
