#include <QCoreApplication>
#include <QDBusConnection>
#include <QDBusError>
#include <QDBusInterface>
#include <QDBusReply>
#include <QDBusUnixFileDescriptor>
#include <QDir>
#include <QFile>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QTemporaryFile>
#include <QTextStream>
#include <QTimer>
#include <QVariantMap>

#include <unistd.h>

class WindowReceiver : public QObject
{
    Q_OBJECT
    Q_CLASSINFO("D-Bus Interface", "org.codex.ReforgerQueueKWinProbe")

public:
    QString json;
    QString error;
    bool done = false;

public slots:
    void WindowsJson(const QString &payload)
    {
        json = payload;
        done = true;
        QCoreApplication::quit();
    }

    void Error(const QString &message)
    {
        error = message;
        done = true;
        QCoreApplication::quit();
    }
};

QDBusConnection sessionBus()
{
    QDBusConnection bus = QDBusConnection::sessionBus();
    if (bus.isConnected()) {
        return bus;
    }

    QString runtimeDir = qEnvironmentVariable("XDG_RUNTIME_DIR");
    if (runtimeDir.isEmpty()) {
        runtimeDir = QStringLiteral("/run/user/%1").arg(getuid());
    }

    return QDBusConnection::connectToBus(
        QStringLiteral("unix:path=%1/bus").arg(runtimeDir),
        QStringLiteral("reforger-queue-session"));
}

QString jsString(const QString &value)
{
    QJsonArray array;
    array.append(value);
    QByteArray encoded = QJsonDocument(array).toJson(QJsonDocument::Compact);
    encoded.chop(1);
    encoded.remove(0, 1);
    return QString::fromUtf8(encoded);
}

int capture(const QString &method, const QString &target, const QString &outputPath, QTextStream &out, QTextStream &err)
{
    QFile file(outputPath);
    if (!file.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
        err << "failed to open output file: " << file.errorString() << "\n";
        return 1;
    }

    QDBusConnection bus = sessionBus();
    if (!bus.isConnected()) {
        err << "failed to connect to the session D-Bus: " << bus.lastError().message() << "\n";
        return 1;
    }

    QDBusInterface iface(
        QStringLiteral("org.kde.KWin"),
        QStringLiteral("/org/kde/KWin/ScreenShot2"),
        QStringLiteral("org.kde.KWin.ScreenShot2"),
        bus);

    if (!iface.isValid()) {
        err << "failed to connect to KWin ScreenShot2: "
            << iface.lastError().message() << "\n";
        return 1;
    }

    QVariantMap options;
    QDBusUnixFileDescriptor fd(file.handle());
    QDBusReply<QVariantMap> reply = iface.call(method, target, options, QVariant::fromValue(fd));

    file.close();

    if (!reply.isValid()) {
        err << method << " failed for " << target << ": "
            << reply.error().name() << ": " << reply.error().message() << "\n";
        return 1;
    }

    QJsonObject resultJson;
    const QVariantMap resultMap = reply.value();
    for (auto it = resultMap.constBegin(); it != resultMap.constEnd(); ++it) {
        resultJson.insert(it.key(), QJsonValue::fromVariant(it.value()));
    }
    out << QJsonDocument(resultJson).toJson(QJsonDocument::Compact) << "\n";

    return 0;
}

QString windowProbeScript(const QString &service, const QString &path, const QString &interface)
{
    return QStringLiteral(R"JS(
function safeString(value) {
    if (value === undefined || value === null) {
        return "";
    }
    return String(value);
}

function safeNumber(value) {
    var number = Number(value);
    return isNaN(number) ? 0 : number;
}

function safeBool(value) {
    return Boolean(value);
}

function prop(object, name, fallback) {
    try {
        if (object[name] === undefined || object[name] === null) {
            return fallback;
        }
        return object[name];
    } catch (error) {
        return fallback;
    }
}

function rect(value) {
    if (!value) {
        return {"x": 0, "y": 0, "width": 0, "height": 0};
    }
    return {
        "x": Math.round(safeNumber(prop(value, "x", 0))),
        "y": Math.round(safeNumber(prop(value, "y", 0))),
        "width": Math.round(safeNumber(prop(value, "width", 0))),
        "height": Math.round(safeNumber(prop(value, "height", 0)))
    };
}

function outputName(window) {
    try {
        if (window.output && window.output.name) {
            return String(window.output.name);
        }
    } catch (error) {
    }
    return "";
}

function windowInfo(window) {
    var geometry = rect(prop(window, "frameGeometry", prop(window, "bufferGeometry", null)));
    return {
        "internalId": safeString(prop(window, "internalId", "")),
        "pid": safeNumber(prop(window, "pid", 0)),
        "caption": safeString(prop(window, "caption", "")),
        "resourceClass": safeString(prop(window, "resourceClass", "")),
        "resourceName": safeString(prop(window, "resourceName", "")),
        "windowRole": safeString(prop(window, "windowRole", "")),
        "desktopFileName": safeString(prop(window, "desktopFileName", "")),
        "output": outputName(window),
        "frameGeometry": geometry,
        "normalWindow": safeBool(prop(window, "normalWindow", false)),
        "fullScreen": safeBool(prop(window, "fullScreen", false)),
        "minimized": safeBool(prop(window, "minimized", false)),
        "hidden": safeBool(prop(window, "hidden", false)),
        "keepAbove": safeBool(prop(window, "keepAbove", false)),
        "skipTaskbar": safeBool(prop(window, "skipTaskbar", false))
    };
}

try {
    var windows = workspace.windowList().map(windowInfo);
    callDBus(%1, %2, %3, "WindowsJson", JSON.stringify(windows));
} catch (error) {
    callDBus(%1, %2, %3, "Error", safeString(error && error.stack ? error.stack : error));
}
)JS")
        .arg(jsString(service), jsString(path), jsString(interface));
}

int listWindows(QCoreApplication &app, QTextStream &out, QTextStream &err)
{
    QDBusConnection bus = sessionBus();
    if (!bus.isConnected()) {
        err << "failed to connect to the session D-Bus: " << bus.lastError().message() << "\n";
        return 1;
    }

    const QString service = QStringLiteral("org.codex.ReforgerQueueKWinProbe%1").arg(QCoreApplication::applicationPid());
    const QString path = QStringLiteral("/Probe");
    const QString interface = QStringLiteral("org.codex.ReforgerQueueKWinProbe");

    WindowReceiver receiver;
    if (!bus.registerService(service)) {
        err << "failed to register D-Bus service " << service << ": " << bus.lastError().message() << "\n";
        return 1;
    }
    if (!bus.registerObject(path, &receiver, QDBusConnection::ExportAllSlots)) {
        err << "failed to register D-Bus object " << path << ": " << bus.lastError().message() << "\n";
        bus.unregisterService(service);
        return 1;
    }

    QTemporaryFile scriptFile;
    scriptFile.setFileTemplate(QDir::tempPath() + QStringLiteral("/reforger-kwin-window-probe-XXXXXX.js"));
    if (!scriptFile.open()) {
        err << "failed to create temporary KWin script: " << scriptFile.errorString() << "\n";
        bus.unregisterObject(path);
        bus.unregisterService(service);
        return 1;
    }
    scriptFile.write(windowProbeScript(service, path, interface).toUtf8());
    scriptFile.close();

    const QString pluginName = QStringLiteral("reforger-queue-window-probe-%1").arg(QCoreApplication::applicationPid());
    QDBusInterface scripting(
        QStringLiteral("org.kde.KWin"),
        QStringLiteral("/Scripting"),
        QStringLiteral("org.kde.kwin.Scripting"),
        bus);

    if (!scripting.isValid()) {
        err << "failed to connect to KWin scripting: " << scripting.lastError().message() << "\n";
        bus.unregisterObject(path);
        bus.unregisterService(service);
        return 1;
    }

    QDBusReply<int> loadReply = scripting.call(QStringLiteral("loadScript"), scriptFile.fileName(), pluginName);
    if (!loadReply.isValid() || loadReply.value() < 0) {
        err << "failed to load KWin window probe script: "
            << loadReply.error().name() << ": " << loadReply.error().message() << "\n";
        bus.unregisterObject(path);
        bus.unregisterService(service);
        return 1;
    }

    const int scriptId = loadReply.value();
    QDBusInterface script(
        QStringLiteral("org.kde.KWin"),
        QStringLiteral("/Scripting/Script%1").arg(scriptId),
        QStringLiteral("org.kde.kwin.Script"),
        bus);
    QDBusMessage runReply = script.call(QStringLiteral("run"));
    if (runReply.type() == QDBusMessage::ErrorMessage) {
        err << "failed to run KWin window probe script: "
            << runReply.errorName() << ": " << runReply.errorMessage() << "\n";
        scripting.call(QStringLiteral("unloadScript"), pluginName);
        bus.unregisterObject(path);
        bus.unregisterService(service);
        return 1;
    }

    QTimer::singleShot(5000, &app, [&receiver]() {
        if (!receiver.done) {
            receiver.error = QStringLiteral("timed out waiting for KWin window probe results");
            receiver.done = true;
            QCoreApplication::quit();
        }
    });
    app.exec();

    scripting.call(QStringLiteral("unloadScript"), pluginName);
    bus.unregisterObject(path);
    bus.unregisterService(service);

    if (!receiver.error.isEmpty()) {
        err << receiver.error << "\n";
        return 1;
    }
    if (receiver.json.isEmpty()) {
        err << "KWin window probe returned no data\n";
        return 1;
    }

    out << receiver.json << "\n";
    return 0;
}

int main(int argc, char *argv[])
{
    QCoreApplication app(argc, argv);
    QTextStream out(stdout);
    QTextStream err(stderr);

    const QStringList args = QCoreApplication::arguments();

    if (args.size() == 3 && !args.at(1).startsWith(QLatin1String("--"))) {
        return capture(QStringLiteral("CaptureScreen"), args.at(1), args.at(2), out, err);
    }

    if (args.size() == 2 && args.at(1) == QLatin1String("--list-windows")) {
        return listWindows(app, out, err);
    }

    if (args.size() == 4 && args.at(1) == QLatin1String("--capture-screen")) {
        return capture(QStringLiteral("CaptureScreen"), args.at(2), args.at(3), out, err);
    }

    if (args.size() == 4 && args.at(1) == QLatin1String("--capture-window")) {
        return capture(QStringLiteral("CaptureWindow"), args.at(2), args.at(3), out, err);
    }

    err << "usage:\n"
        << "  " << args.value(0) << " OUTPUT_NAME OUTPUT_FILE\n"
        << "  " << args.value(0) << " --capture-screen OUTPUT_NAME OUTPUT_FILE\n"
        << "  " << args.value(0) << " --capture-window INTERNAL_ID OUTPUT_FILE\n"
        << "  " << args.value(0) << " --list-windows\n";
    return 64;
}

#include "kwin_capture_screen.moc"