#include "adblock.h"

#include "att_adblock.h"

#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QFutureWatcher>
#include <QNetworkAccessManager>
#include <QNetworkReply>
#include <QSaveFile>
#include <QStandardPaths>
#include <QWebEngineUrlRequestInfo>
#include <QtConcurrent>

namespace {

// qutebrowser's default content.blocking.adblock.lists.
constexpr const char *kLists[] = {
    "https://easylist.to/easylist/easylist.txt",
    "https://easylist.to/easylist/easyprivacy.txt",
};
constexpr int kListCount = std::size(kLists);

using EnginePtr = std::shared_ptr<const AttAdblockEngine>;

EnginePtr wrap(AttAdblockEngine *engine)
{
    if (!engine)
        return {};
    return EnginePtr(engine, [](const AttAdblockEngine *e) {
        att_adblock_engine_free(const_cast<AttAdblockEngine *>(e));
    });
}

EnginePtr loadCache(const QString &path)
{
    QFile file(path);
    if (!file.open(QIODevice::ReadOnly))
        return {};
    const QByteArray data = file.readAll();
    return wrap(att_adblock_engine_deserialize(
        reinterpret_cast<const uint8_t *>(data.constData()), size_t(data.size())));
}

bool saveCache(const AttAdblockEngine *engine, const QString &path)
{
    size_t len = 0;
    uint8_t *data = att_adblock_engine_serialize(engine, &len);
    if (!data)
        return false;

    QDir().mkpath(QFileInfo(path).absolutePath());
    QSaveFile file(path);
    const bool ok = file.open(QIODevice::WriteOnly)
                    && file.write(reinterpret_cast<const char *>(data), qint64(len)) == qint64(len)
                    && file.commit();
    att_adblock_buffer_free(data, len);
    return ok;
}

// The webRequest names adblock-rust matches $script, $image, ... options
// against. Same mapping as qutebrowser's braveadblock.py.
const char *requestType(QWebEngineUrlRequestInfo::ResourceType type)
{
    switch (type) {
    case QWebEngineUrlRequestInfo::ResourceTypeMainFrame:
        return "main_frame";
    case QWebEngineUrlRequestInfo::ResourceTypeSubFrame:
        return "sub_frame";
    case QWebEngineUrlRequestInfo::ResourceTypeStylesheet:
        return "stylesheet";
    case QWebEngineUrlRequestInfo::ResourceTypeScript:
        return "script";
    case QWebEngineUrlRequestInfo::ResourceTypeImage:
    case QWebEngineUrlRequestInfo::ResourceTypeFavicon:
        return "image";
    case QWebEngineUrlRequestInfo::ResourceTypeFontResource:
        return "font";
    case QWebEngineUrlRequestInfo::ResourceTypeObject:
        return "object";
    case QWebEngineUrlRequestInfo::ResourceTypeMedia:
        return "media";
    case QWebEngineUrlRequestInfo::ResourceTypeXhr:
        return "xhr";
    case QWebEngineUrlRequestInfo::ResourceTypePing:
        return "ping";
    case QWebEngineUrlRequestInfo::ResourceTypeCspReport:
        return "csp_report";
    case QWebEngineUrlRequestInfo::ResourceTypeWebSocket:
        return "websocket";
    default:
        return "other";
    }
}

} // namespace

struct AdBlocker::Download
{
    QByteArray lists[kListCount];
    int remaining = kListCount;
    int failed = 0;
};

AdBlocker::AdBlocker(QObject *parent)
    : QWebEngineUrlRequestInterceptor(parent)
    , m_network(new QNetworkAccessManager(this))
    , m_cachePath(QStandardPaths::writableLocation(QStandardPaths::AppDataLocation)
                  + QStringLiteral("/adblock-cache.dat"))
{
}

void AdBlocker::interceptRequest(QWebEngineUrlRequestInfo &info)
{
    if (!m_enabled)
        return;

    const QUrl url = info.requestUrl();
    const QString scheme = url.scheme();
    if (scheme != QLatin1String("https") && scheme != QLatin1String("http")
        && scheme != QLatin1String("wss") && scheme != QLatin1String("ws"))
        return;

    const EnginePtr current = engine();
    if (!current)
        return;

    const QByteArray requestUrl = url.toEncoded();
    const QByteArray sourceUrl = info.firstPartyUrl().toEncoded();
    const QByteArray method = info.requestMethod();

    if (att_adblock_should_block(current.get(), requestUrl.constData(), sourceUrl.constData(),
                                 requestType(info.resourceType()), method.constData()))
        info.block(true);
}

void AdBlocker::start()
{
    auto *watcher = new QFutureWatcher<EnginePtr>(this);
    connect(watcher, &QFutureWatcher<EnginePtr>::finished, this, [this, watcher] {
        watcher->deleteLater();
        if (EnginePtr cached = watcher->result())
            setEngine(std::move(cached));
        else
            update(); // First run, or a cache from an incompatible adblock-rust.
    });
    watcher->setFuture(QtConcurrent::run(loadCache, m_cachePath));
}

void AdBlocker::update()
{
    if (m_updating) {
        emit message(QStringLiteral("Adblock update already running"), false);
        return;
    }
    m_updating = true;
    emit message(QStringLiteral("Downloading adblock lists..."), false);

    auto download = std::make_shared<Download>();
    for (int i = 0; i < kListCount; ++i) {
        QNetworkReply *reply = m_network->get(QNetworkRequest(QUrl(QString::fromLatin1(kLists[i]))));
        connect(reply, &QNetworkReply::finished, this, [this, reply, download, i] {
            reply->deleteLater();
            if (reply->error() == QNetworkReply::NoError)
                download->lists[i] = reply->readAll();
            else
                download->failed++;
            if (--download->remaining == 0)
                compile(download);
        });
    }
}

void AdBlocker::compile(const std::shared_ptr<Download> &download)
{
    if (download->failed == kListCount) {
        m_updating = false;
        emit message(QStringLiteral("Adblock update failed: could not download any list"), true);
        return;
    }

    QByteArray rules;
    for (const QByteArray &list : download->lists)
        rules += list + '\n';

    // Compiling EasyList takes a moment, so keep it off the GUI thread.
    auto *watcher = new QFutureWatcher<EnginePtr>(this);
    connect(watcher, &QFutureWatcher<EnginePtr>::finished, this, [this, watcher, download] {
        watcher->deleteLater();
        m_updating = false;
        EnginePtr compiled = watcher->result();
        if (!compiled) {
            emit message(QStringLiteral("Adblock update failed: could not compile the lists"), true);
            return;
        }
        setEngine(std::move(compiled));
        if (download->failed > 0)
            emit message(QStringLiteral("Adblock updated, but %1 of %2 lists failed to download")
                             .arg(download->failed)
                             .arg(kListCount),
                         true);
        else
            emit message(QStringLiteral("Adblock lists updated"), false);
    });
    watcher->setFuture(QtConcurrent::run([rules, path = m_cachePath] {
        EnginePtr engine = wrap(att_adblock_engine_from_rules(
            reinterpret_cast<const uint8_t *>(rules.constData()), size_t(rules.size())));
        if (engine)
            saveCache(engine.get(), path);
        return engine;
    }));
}

void AdBlocker::setEngine(EnginePtr engine)
{
    const std::lock_guard lock(m_engineMutex);
    m_engine = std::move(engine);
}

AdBlocker::EnginePtr AdBlocker::engine() const
{
    const std::lock_guard lock(m_engineMutex);
    return m_engine;
}