#include "SyncSession.h"

#include <QHttpMultiPart>
#include <QJsonValue>
#include <QNetworkReply>
#include <QNetworkRequest>
#include <QTemporaryFile>
#include <QWebSocket>
#include <QWebSocketProtocol>

using namespace Qt::StringLiterals;

SyncSession::SyncSession(QWebSocket *socket, QObject *parent) :
    QObject{parent},
    m_socket{socket}
{
    m_requestTimer.setSingleShot(true);
    m_requestTimer.setInterval(10'000);
    connect(&m_requestTimer, &QTimer::timeout, this, [this] {
        send({{u"type"_s, u"error"_s}, {u"message"_s, u"No job request received within 10 seconds"_s}});
        finishJob(u"Job request timeout"_s);
    });
    m_requestTimer.start();
    send({
        {u"type"_s, u"hello"_s},
        {u"protocol"_s, 2},
        {u"message"_s, u"Send one complete inspection or synchronization job"_s},
    });
}

SyncSession::~SyncSession()
{
    abort();
}

void SyncSession::handleTextMessage(const QString &message)
{
    QJsonParseError parseError;
    const QJsonValue rootValue = QJsonValue::fromJson(message.toUtf8(), &parseError);
    if (parseError.error != QJsonParseError::NoError || !rootValue.isObject())
    {
        send({{u"type"_s, u"error"_s}, {u"message"_s, u"Invalid WebSocket JSON request"_s}});
        finishJob(u"Invalid job request"_s);
        return;
    }

    const QJsonObject request = rootValue.toObject();
    const QString type = request.value(u"type"_s).toString();
    if (type == u"abort")
    {
        abort();
        finishJob(u"Job aborted"_s);
        return;
    }
    if (m_jobKind != JobKind::None)
    {
        send({{u"type"_s, u"error"_s}, {u"message"_s, u"Each WebSocket accepts exactly one job"_s}});
        abort();
        finishJob(u"Unexpected additional request"_s);
        return;
    }

    if (type == u"inspect")
    {
        m_requestTimer.stop();
        m_jobKind = JobKind::Inspection;
        inspect(request);
    }
    else if (type == u"sync")
    {
        const QString direction = request.value(u"direction"_s).toString();
        if (direction != u"left-to-right" && direction != u"right-to-left" && direction != u"bidirectional")
        {
            send({{u"type"_s, u"error"_s}, {u"message"_s, u"Unknown synchronization direction"_s}});
            finishJob(u"Invalid job request"_s);
            return;
        }
        m_requestTimer.stop();
        m_jobKind = JobKind::Synchronization;
        m_direction = direction;
        inspect(request);
    }
    else
    {
        send({{u"type"_s, u"error"_s}, {u"message"_s, u"Unknown request type"_s}});
        finishJob(u"Invalid job request"_s);
    }
}

void SyncSession::inspect(const QJsonObject &request)
{
    reset();
    send({{u"type"_s, u"inspection-status"_s},
          {u"stage"_s, u"starting"_s},
          {u"message"_s, u"Inspecting both shared links"_s}});
    configureSide(u"left"_s, request.value(u"left"_s));
    configureSide(u"right"_s, request.value(u"right"_s));
}

void SyncSession::configureSide(const QString &side, const QJsonValue &value)
{
    if (!value.isObject())
    {
        sendSideStatus(side, u"error"_s, u"Missing share configuration"_s);
        sideResolved(side);
        return;
    }
    const QJsonObject config = value.toObject();
    QString parseError;
    ShareEndpoint endpoint = ShareEndpoint::fromUserInput(config.value(u"url"_s).toString(),
                                                          config.value(u"password"_s).toString(), &parseError);
    if (!parseError.isEmpty())
    {
        sendSideStatus(side, u"error"_s, parseError);
        sideResolved(side);
        return;
    }

    auto *client = new ImmichClient{std::move(endpoint), this};
    if (side == u"left")
    {
        m_leftClient = client;
    }
    else
    {
        m_rightClient = client;
    }

    connect(client, &ImmichClient::inspectionStage, this,
            [this, side](const QString &stage, const QString &text) { sendSideStatus(side, stage, text); });
    connect(client, &ImmichClient::inspected, this, [this, side](const ShareInfo &info) { sideInspected(side, info); });
    connect(client, &ImmichClient::passwordRequired, this, [this, side] {
        sendSideStatus(side, u"password-required"_s, u"This share needs a password"_s);
        sideResolved(side);
    });
    connect(client, &ImmichClient::failed, this, [this, side](const QString &message) {
        sendSideStatus(side, u"error"_s, message);
        sideResolved(side);
    });
    client->inspect();
}

void SyncSession::sideInspected(const QString &side, const ShareInfo &info)
{
    if (side == u"left")
    {
        m_leftInfo = info;
        m_leftUsable = true;
    }
    else
    {
        m_rightInfo = info;
        m_rightUsable = true;
    }
    const QJsonObject details =
        shareInfoJson(side == u"left" ? m_leftClient->endpoint() : m_rightClient->endpoint(), info);
    QJsonObject status{{u"type"_s, u"side-status"_s},
                       {u"side"_s, side},
                       {u"stage"_s, u"ready"_s},
                       {u"message"_s, u"Share is ready"_s},
                       {u"share"_s, details}};
    send(std::move(status));
    sideResolved(side);
}

void SyncSession::sideResolved(const QString &side)
{
    if (side == u"left")
    {
        if (m_leftResolved)
        {
            return;
        }
        m_leftResolved = true;
    }
    else
    {
        if (m_rightResolved)
        {
            return;
        }
        m_rightResolved = true;
    }
    if (m_leftResolved && m_rightResolved)
    {
        sendInspection();
    }
}

bool SyncSession::directionAllowed(const QString &direction, QString *reason) const
{
    bool allowed = false;
    QString why;
    if (!m_ready)
    {
        why = u"Both links must be inspected successfully"_s;
    }
    else if (direction == u"left-to-right")
    {
        allowed = m_leftInfo.allowDownload && m_rightInfo.allowUpload;
        if (!m_leftInfo.allowDownload)
        {
            why = u"The left share does not allow downloads"_s;
        }
        else if (!m_rightInfo.allowUpload)
        {
            why = u"The right share does not allow uploads"_s;
        }
    }
    else if (direction == u"right-to-left")
    {
        allowed = m_rightInfo.allowDownload && m_leftInfo.allowUpload;
        if (!m_rightInfo.allowDownload)
        {
            why = u"The right share does not allow downloads"_s;
        }
        else if (!m_leftInfo.allowUpload)
        {
            why = u"The left share does not allow uploads"_s;
        }
    }
    else if (direction == u"bidirectional")
    {
        QString leftReason;
        QString rightReason;
        allowed =
            directionAllowed(u"left-to-right"_s, &leftReason) && directionAllowed(u"right-to-left"_s, &rightReason);
        if (!allowed)
        {
            why = leftReason.isEmpty() ? rightReason : leftReason;
        }
    }
    else
    {
        why = u"Unknown synchronization direction"_s;
    }
    if (reason)
    {
        *reason = why;
    }
    return allowed;
}

void SyncSession::sendInspection()
{
    m_ready = m_leftUsable && m_rightUsable;
    if (!m_ready)
    {
        send({{u"type"_s, u"inspection"_s},
              {u"ready"_s, false},
              {u"message"_s, u"Both links must be ready before synchronization can start"_s}});
        finishJob(u"Inspection failed"_s);
        return;
    }

    m_plan = buildSyncPlan(m_leftInfo, m_rightInfo);
    const auto option = [this](const QString &direction, qsizetype count) {
        QString reason;
        const bool allowed = directionAllowed(direction, &reason);
        return QJsonObject{{u"allowed"_s, allowed}, {u"missing"_s, static_cast<qint64>(count)}, {u"reason"_s, reason}};
    };

    const QJsonObject options{
        {u"left-to-right"_s, option(u"left-to-right"_s, m_plan.leftToRight.size())},
        {u"right-to-left"_s, option(u"right-to-left"_s, m_plan.rightToLeft.size())},
        {u"bidirectional"_s, option(u"bidirectional"_s, m_plan.leftToRight.size() + m_plan.rightToLeft.size())},
    };
    send({
        {u"type"_s, u"inspection"_s},
        {u"ready"_s, true},
        {u"left"_s, shareInfoJson(m_leftClient->endpoint(), m_leftInfo)},
        {u"right"_s, shareInfoJson(m_rightClient->endpoint(), m_rightInfo)},
        {u"options"_s, options},
        {u"message"_s, u"Inspection complete"_s},
    });
    if (m_jobKind == JobKind::Inspection)
    {
        finishJob(u"Inspection complete"_s);
        return;
    }
    startSync(m_direction);
}

void SyncSession::startSync(const QString &direction)
{
    if (m_active)
    {
        send({{u"type"_s, u"error"_s}, {u"message"_s, u"A synchronization is already running"_s}});
        return;
    }
    QString reason;
    if (!directionAllowed(direction, &reason))
    {
        send({{u"type"_s, u"error"_s}, {u"message"_s, reason}});
        finishJob(u"Synchronization unavailable"_s);
        return;
    }

    m_queue.clear();
    if (direction == u"left-to-right" || direction == u"bidirectional")
    {
        for (const AssetInfo &asset : m_plan.leftToRight)
        {
            m_queue.append({m_leftClient, m_rightClient, asset, u"left-to-right"_s});
        }
    }
    if (direction == u"right-to-left" || direction == u"bidirectional")
    {
        for (const AssetInfo &asset : m_plan.rightToLeft)
        {
            m_queue.append({m_rightClient, m_leftClient, asset, u"right-to-left"_s});
        }
    }
    m_total = m_queue.size();
    m_completed = 0;
    m_failed = 0;
    m_active = true;
    qInfo().noquote() << u"Synchronization started: %1, %2 asset(s)"_s.arg(direction).arg(m_total);
    send({{u"type"_s, u"sync-status"_s},
          {u"stage"_s, u"started"_s},
          {u"direction"_s, direction},
          {u"total"_s, static_cast<qint64>(m_total)},
          {u"message"_s, u"Synchronization started"_s}});
    startNextTransfer();
}

void SyncSession::startNextTransfer()
{
    if (!m_active)
    {
        return;
    }
    if (m_queue.isEmpty())
    {
        m_active = false;
        send({{u"type"_s, u"sync-status"_s},
              {u"stage"_s, u"complete"_s},
              {u"completed"_s, static_cast<qint64>(m_completed)},
              {u"failed"_s, static_cast<qint64>(m_failed)},
              {u"total"_s, static_cast<qint64>(m_total)},
              {u"message"_s,
               m_failed == 0 ? u"Synchronization complete"_s : u"Synchronization complete with errors"_s}});
        finishJob(u"Synchronization complete"_s);
        return;
    }

    const TransferSpec &transfer = m_queue.front();
    const quint64 serial = ++m_transferSerial;
    qInfo().noquote() << u"Transfer %1/%2 started: %3 (%4)"_s
                             .arg(m_completed + m_failed + 1)
                             .arg(m_total)
                             .arg(transfer.asset.fileName, transfer.direction);
    send({{u"type"_s, u"asset-status"_s},
          {u"stage"_s, u"downloading"_s},
          {u"direction"_s, transfer.direction},
          {u"fileName"_s, transfer.asset.fileName},
          {u"current"_s, static_cast<qint64>(m_completed + m_failed + 1)},
          {u"total"_s, static_cast<qint64>(m_total)},
          {u"message"_s, u"Downloading original"_s}});

    m_progressTimer.restart();
    m_temporaryFile = new QTemporaryFile{u"/tmp/immich-sync-XXXXXX"_s, this};
    if (!m_temporaryFile->open())
    {
        finishTransfer(serial, false, u"Unable to create a temporary transfer file"_s);
        return;
    }
    m_download = transfer.source->download(transfer.asset);
    connect(m_download, &QNetworkReply::readyRead, this, [this, serial] { storeDownloadedData(serial); });
    connect(m_download, &QNetworkReply::downloadProgress, this, [this, serial](qint64 received, qint64 total) {
        if (serial == m_transferSerial)
        {
            sendProgress(u"downloading"_s, received, total);
        }
    });
    connect(m_download, &QNetworkReply::finished, this, [this, serial] {
        if (serial != m_transferSerial || !m_download)
        {
            return;
        }
        if (!storeDownloadedData(serial))
        {
            return;
        }
        const int status = m_download->attribute(QNetworkRequest::HttpStatusCodeAttribute).toInt();
        if (m_download->error() != QNetworkReply::NoError || status < 200 || status >= 300)
        {
            finishTransfer(serial, false, u"Download failed: %1"_s.arg(m_download->errorString()));
            return;
        }
        if (!m_temporaryFile->flush() || !m_temporaryFile->seek(0))
        {
            finishTransfer(serial, false, u"Unable to rewind the temporary transfer file"_s);
            return;
        }
        beginUpload(serial);
    });
}

bool SyncSession::storeDownloadedData(quint64 serial)
{
    if (serial != m_transferSerial || !m_download || !m_temporaryFile)
    {
        return false;
    }
    while (m_download->bytesAvailable() > 0)
    {
        const QByteArray chunk = m_download->read(1024 * 1024);
        if (chunk.isEmpty())
        {
            break;
        }
        if (m_temporaryFile->write(chunk) != chunk.size())
        {
            finishTransfer(serial, false, u"Unable to write the temporary transfer file"_s);
            return false;
        }
    }
    return true;
}

void SyncSession::beginUpload(quint64 serial)
{
    if (serial != m_transferSerial || !m_download || !m_temporaryFile || m_upload)
    {
        return;
    }
    const int status = m_download->attribute(QNetworkRequest::HttpStatusCodeAttribute).toInt();
    if (status < 200 || status >= 300)
    {
        return;
    }
    const TransferSpec &transfer = m_queue.front();
    m_download->deleteLater();
    m_download = nullptr;
    m_multipart = new QHttpMultiPart{QHttpMultiPart::FormDataType};
    m_upload = transfer.destination->upload(transfer.asset, m_temporaryFile, m_multipart);
    m_multipart->setParent(m_upload);
    m_uploadCompleteNotified = false;
    send({{u"type"_s, u"asset-status"_s},
          {u"stage"_s, u"uploading"_s},
          {u"direction"_s, transfer.direction},
          {u"fileName"_s, transfer.asset.fileName},
          {u"current"_s, static_cast<qint64>(m_completed + m_failed + 1)},
          {u"total"_s, static_cast<qint64>(m_total)},
          {u"message"_s, u"Uploading original to destination"_s}});
    connect(m_upload, &QNetworkReply::uploadProgress, this, [this, serial](qint64 sent, qint64 total) {
        if (serial == m_transferSerial)
        {
            sendProgress(u"uploading"_s, sent, total);
            if (total > 0 && sent >= total && !m_uploadCompleteNotified && !m_queue.isEmpty())
            {
                m_uploadCompleteNotified = true;
                const TransferSpec &currentTransfer = m_queue.front();
                send({{u"type"_s, u"asset-status"_s},
                      {u"stage"_s, u"processing"_s},
                      {u"direction"_s, currentTransfer.direction},
                      {u"fileName"_s, currentTransfer.asset.fileName},
                      {u"current"_s, static_cast<qint64>(m_completed + m_failed + 1)},
                      {u"total"_s, static_cast<qint64>(m_total)},
                      {u"message"_s, u"Destination Immich is processing the asset"_s}});
            }
        }
    });
    connect(m_upload, &QNetworkReply::finished, this, [this, serial] {
        if (serial != m_transferSerial || !m_upload)
        {
            return;
        }
        const QByteArray body = m_upload->readAll();
        const int uploadStatus = m_upload->attribute(QNetworkRequest::HttpStatusCodeAttribute).toInt();
        if (m_upload->error() != QNetworkReply::NoError || (uploadStatus != 200 && uploadStatus != 201))
        {
            finishTransfer(serial, false, u"Upload failed: %1"_s.arg(m_upload->errorString()));
            return;
        }
        const QJsonValue rootValue = QJsonValue::fromJson(body);
        const QString result = rootValue.toObject().value(u"status"_s).toString();
        finishTransfer(serial, true,
                       uploadStatus == 200 ? u"Existing destination asset added to share"_s : u"Asset synchronized"_s,
                       result);
    });
}

void SyncSession::finishTransfer(quint64 serial, bool success, const QString &message, const QString &result)
{
    if (serial != m_transferSerial || m_queue.isEmpty())
    {
        return;
    }
    ++m_transferSerial;
    const TransferSpec transfer = m_queue.takeFirst();
    qInfo().noquote() << u"Transfer finished: %1 (%2): %3"_s.arg(transfer.asset.fileName,
                                                                 success ? u"success"_s : u"failed"_s, message);
    if (success)
    {
        ++m_completed;
    }
    else
    {
        ++m_failed;
        if (m_upload)
        {
            m_upload->abort();
        }
        if (m_download)
        {
            m_download->abort();
        }
    }
    send({{u"type"_s, u"asset-status"_s},
          {u"stage"_s, success ? u"complete"_s : u"failed"_s},
          {u"direction"_s, transfer.direction},
          {u"fileName"_s, transfer.asset.fileName},
          {u"result"_s, result},
          {u"completed"_s, static_cast<qint64>(m_completed)},
          {u"failed"_s, static_cast<qint64>(m_failed)},
          {u"total"_s, static_cast<qint64>(m_total)},
          {u"message"_s, message}});

    if (m_upload)
    {
        m_upload->deleteLater();
    }
    if (m_download)
    {
        m_download->deleteLater();
    }
    m_upload = nullptr;
    m_download = nullptr;
    m_multipart = nullptr;
    m_uploadCompleteNotified = false;
    if (m_temporaryFile)
    {
        m_temporaryFile->deleteLater();
        m_temporaryFile = nullptr;
    }
    startNextTransfer();
}

void SyncSession::sendProgress(const QString &phase, qint64 transferred, qint64 total)
{
    if (m_progressTimer.isValid() && m_progressTimer.elapsed() < 200 && transferred != total)
    {
        return;
    }
    m_progressTimer.restart();
    if (m_queue.isEmpty())
    {
        return;
    }
    const TransferSpec &transfer = m_queue.front();
    send({{u"type"_s, u"asset-progress"_s},
          {u"phase"_s, phase},
          {u"direction"_s, transfer.direction},
          {u"fileName"_s, transfer.asset.fileName},
          {u"bytes"_s, transferred},
          {u"totalBytes"_s, total},
          {u"current"_s, static_cast<qint64>(m_completed + m_failed + 1)},
          {u"total"_s, static_cast<qint64>(m_total)}});
}

void SyncSession::abort()
{
    const bool wasActive = m_active;
    m_active = false;
    ++m_transferSerial;
    m_queue.clear();
    if (m_upload)
    {
        m_upload->abort();
        m_upload->deleteLater();
    }
    if (m_download)
    {
        m_download->abort();
        m_download->deleteLater();
    }
    m_upload = nullptr;
    m_download = nullptr;
    m_multipart = nullptr;
    m_uploadCompleteNotified = false;
    if (m_temporaryFile)
    {
        m_temporaryFile->deleteLater();
        m_temporaryFile = nullptr;
    }
    if (wasActive)
    {
        qInfo() << "Synchronization aborted because its WebSocket disconnected";
        send({{u"type"_s, u"sync-status"_s},
              {u"stage"_s, u"aborted"_s},
              {u"message"_s, u"Synchronization and network transfers were aborted"_s}});
    }
}

void SyncSession::finishJob(const QString &reason)
{
    m_requestTimer.stop();
    if (!m_socket || m_socket->state() != QAbstractSocket::ConnectedState)
    {
        return;
    }
    const QPointer<QWebSocket> socket = m_socket;
    m_socket->flush();
    m_socket->close(QWebSocketProtocol::CloseCodeNormal, reason);
    QTimer::singleShot(2'000, m_socket, [socket] {
        if (socket && socket->state() != QAbstractSocket::UnconnectedState)
        {
            socket->abort();
        }
    });
}

void SyncSession::reset()
{
    abort();
    if (m_leftClient)
    {
        delete m_leftClient;
    }
    if (m_rightClient)
    {
        delete m_rightClient;
    }
    m_leftClient = nullptr;
    m_rightClient = nullptr;
    m_leftInfo = {};
    m_rightInfo = {};
    m_leftResolved = false;
    m_rightResolved = false;
    m_leftUsable = false;
    m_rightUsable = false;
    m_ready = false;
    m_plan = {};
    m_total = 0;
    m_completed = 0;
    m_failed = 0;
}

void SyncSession::send(QJsonObject message)
{
    if (!m_socket || m_socket->state() != QAbstractSocket::ConnectedState)
    {
        return;
    }
    message.insert(u"timestamp"_s, QDateTime::currentDateTimeUtc().toString(Qt::ISODateWithMs));
    m_socket->sendTextMessage(QString::fromUtf8(QJsonValue{message}.toJson(QJsonValue::JsonFormat::Compact)));
}

void SyncSession::sendSideStatus(const QString &side, const QString &stage, const QString &message)
{
    send({{u"type"_s, u"side-status"_s}, {u"side"_s, side}, {u"stage"_s, stage}, {u"message"_s, message}});
}