#include "loginrpc.h"
#include "cvt.h"

#include <QtNetwork/QTcpServer>
#include <QtNetwork/QTcpSocket>
#include <QDebug>
#include <QSize>
#include <QProcess>
#include <QRegularExpression>

static int doFbOnly(int x, int y, const QStringList &allOutputs);

static int setMode(const QString &virtOut, const QString &name, const QStringList &allOutputs);

static int sizeDiff(const QSize &a, const QSize &b)
{
    return qAbs((a.width() * a.width()) - (b.width() * b.width()))
            + qAbs((a.height() * a.height()) - (b.height() * b.height()));
}

LoginRpc::LoginRpc(int port, QObject *parent)
    : QObject(parent)
{
    QTcpServer *srv = new QTcpServer(this);
    srv->listen(QHostAddress::Any, port);
    connect(srv, &QTcpServer::newConnection, [=] {
        while (srv->hasPendingConnections()) {
            QTcpSocket *sock = srv->nextPendingConnection();
            handleIncoming(sock);
        }
    });
}

void LoginRpc::handleIncoming(QTcpSocket *sock) {
    connect(sock, &QTcpSocket::readyRead, [=] {
        QByteArray ba = sock->readAll(); // XXX We assume everything arrives in one packet
        if (ba.length() < 2) {
            sock->deleteLater();
            return;
        }
        int vers = (ba[0] << 8) + ba[1];
        ba = ba.mid(2);
        ba = QByteArray::fromBase64(ba);
        if (vers == 1) {
            handleCommandV1(QString::fromUtf8(ba));
        } else {
            qDebug() << "Ignoring unknown Login RPC version" << vers;
        }
    });
}

void LoginRpc::handleCommandV1(const QString &command)
{
    QStringList lines = command.split('\n');
    while (lines.count() < 3) {
        lines.append(QString());
    }
    QString res = lines[2];
    QStringList parts = res.split("x");
    qDebug() << "Got resolution" << res << "parsed to" << parts;
    if (parts.size() == 2) {
        struct {
            QString output;
            QString modeName;
            QSize size;
        } bestConMode, bestDisconMode, *currentReadMode = nullptr;
        int x = parts[0].toInt();
        int y = parts[1].toInt();
        qDebug() << "As int:" << x << y;
        if (x > 0 && y > 0) {
            x = (x / 8) * 8;
            y = (y / 8) * 8;
            // TODO: Configurable min max sizes
            if (x < 1024) x = 1024;
            if (y < 720) y = 720;
            if (x > 1920) x = 1920;
            if (y > 1080) y = 1080;
            QString name = QString("%1x%2_60").arg(x).arg(y);
            mode *mode = vert_refresh(x, y, 60, 0, 0, 0);
            QProcess p;
            // Fetch xrandr default output once
            p.setProcessChannelMode(QProcess::MergedChannels);
            p.start("xrandr", QStringList());
            p.waitForFinished(2000);
            p.kill();
            QString xrandrOutput = QString::fromLocal8Bit(p.readAll());
            qDebug() << "Creating new mode via xrandr";
            p.setProcessChannelMode(QProcess::ForwardedChannels);
            QStringList newmode = QStringList() << "--verbose" << "--newmode" << name << QString::asprintf("%.2f", mode->pclk)
                                     << QString::number(mode->hr) << QString::number(mode->hss) << QString::number(mode->hse) << QString::number(mode->hfl)
                                     << QString::number(mode->vr) << QString::number(mode->vss) << QString::number(mode->vse) << QString::number(mode->vfl)
                                     << "-hsync" << "+vsync";
            // Create mode
            p.start("xrandr", newmode);
            p.waitForFinished(2000);
            qDebug() << "Exit code" << p.exitCode();
            bool lowResFallback = p.exitCode() != 0;
            p.kill();
            do {
                // Find best match among existing modes in case --newmode failed, or --addmode fails later
                QSize wantedSize(x, y);
                QStringList lines = xrandrOutput.split(QRegularExpression("[\r\n]+"));
                QString curScreen;
                bool curConnected = false;
                QRegularExpression scr("^([A-Za-z0-9_-]+)\\s+(connected|disconnected)\\s+");
                QRegularExpression mode("^\\s+([0-9]+)x([0-9]+)(\\S*)\\s+\\d");
                for (const auto &line : lines) {
                    //qDebug() << "Line" << line;
                    auto m = scr.match(line);
                    if (m.hasMatch()) {
                        // Is a line that starts a new screen section
                        curScreen = m.captured(1);
                        curConnected = m.captured(2) == QStringLiteral("connected");
                        currentReadMode = curConnected ? &bestConMode : &bestDisconMode;
                        //qDebug() << "Output" << curScreen << curConnected;
                        continue;
                    }
                    if (currentReadMode != nullptr) {
                        m = mode.match(line);
                        if (m.hasMatch()) {
                            // Is a resolution/mode
                            QSize s(m.captured(1).toInt(), m.captured(2).toInt());
                            //qDebug() << "Matched resolution" << m.captured() << s;
                            //qDebug() << "Current:" << sizeDiff(currentReadMode->size, wantedSize)  << ", This: " << sizeDiff(s, wantedSize);
                            if (s.width() > 1000 && s.height() > 700 && sizeDiff(currentReadMode->size, wantedSize) > sizeDiff(s, wantedSize)) {
                                // Better
                                qDebug() << s << "is better than" << currentReadMode->size << "on" << curScreen << curConnected;
                                currentReadMode->output = curScreen;
                                currentReadMode->modeName = m.captured(1) + "x" + m.captured(2) + m.captured(3);
                                currentReadMode->size = s;
                            }
                        }
                    }
                }
            } while(0);
            // Get all outputs
            QRegularExpression re("^(\\S+)\\s.*?(\\w*connected)", QRegularExpression::MultilineOption);
            QRegularExpressionMatchIterator it = re.globalMatch(xrandrOutput);
            QString virtOut, evdiOut;
            QStringList allOutputs, disconnectedOutputs, connectedOutputs;
            while (it.hasNext()) {
                QRegularExpressionMatch m = it.next();
                QString output = m.captured(1);
                if (virtOut.isEmpty() && (output == QLatin1String("VIRTUAL1") || output == QLatin1String("VIRTUAL-1")
                        || output == QLatin1String("Virtual1") || output == QLatin1String("Virtual-1")
                        || output == QLatin1String("default"))) {
                    virtOut = output;
                } else if (output == QLatin1String("DVI-I-1-1")) {
                    evdiOut = output;
                } else {
                    allOutputs << output;
                }
                if (m.captured(2) == QStringLiteral("disconnected")) {
                    disconnectedOutputs << output;
                } else {
                    connectedOutputs << output;
                }
            }
            if (virtOut.isEmpty()) {
                virtOut = evdiOut;
            } else if (!evdiOut.isEmpty()) {
                allOutputs << evdiOut;
            }
            qDebug() << "Virtual output:" << virtOut << "unwanted additional outputs:" << allOutputs;
            int ret = -1;
            // Always try virtual output first, even if --newmode failed. Might have to be re-evaluated in the future
            if (!virtOut.isEmpty()) {
                // Add to virtual output
                p.start("xrandr", QStringList() << "--verbose" << "--addmode" << virtOut << name);
                p.waitForFinished(2000);
                p.kill();
                ret = setMode(virtOut, name, allOutputs);
                if (ret != 0) {
                    allOutputs << virtOut;
                }
            }
            // Overridden, because --newmode failed?
            if (lowResFallback) {
                if (ret != 0 && !bestDisconMode.modeName.isEmpty()) {
                    qDebug() << "Ret" << ret << "- Trying to enable best disconnected screen (1)";
                    ret = setMode(bestDisconMode.output, bestDisconMode.modeName, allOutputs);
                    bestDisconMode.modeName.clear();
                }
                if (ret != 0 && !bestConMode.modeName.isEmpty()) {
                    qDebug() << "Ret" << ret << "- Trying to enable best connected screen (1)";
                    ret = setMode(bestConMode.output, bestConMode.modeName, allOutputs);
                    bestConMode.modeName.clear();
                }
            } else {
                if (ret != 0) {
                    // Do the --fb thing - we had that before and there was a problem, but I don't remember what it
                    // was, and initial testing seems fine!? Not taking notes FTW!
                    ret = doFbOnly(x, y, allOutputs);
                }
            }
            // Either -1 if we didn't have a virtual one, or != 0 if xrandr setting failed
            // Now as fallback, try enabling a disconnected output only
            if (ret != 0 && !disconnectedOutputs.isEmpty()) {
                qDebug() << "Ret:" << ret << "- Trying to enable one random disconnected output";
                QRegularExpression re("\\d+$");
                QString dis;
                int used[10] = {};
                for (const auto & c : connectedOutputs) {
                    auto m = re.match(c);
                    if (m.hasMatch()) {
                        int i = m.captured().toInt();
                        if (i >= 0 && i < 10) {
                            used[i]++;
                        }
                    }
                }
                for (const auto & c : disconnectedOutputs) {
                    auto m = re.match(c);
                    if (m.hasMatch()) {
                        int i = m.captured().toInt();
                        if (i >= 0 && i < 10 && used[i] == 0) {
                            dis = c;
                            break;
                        }
                    }
                }
                if (dis.isEmpty()) {
                    dis = disconnectedOutputs.first();
                }
                // Add to output
                p.start("xrandr", QStringList() << "--verbose" << "--addmode" << dis << name);
                p.waitForFinished(2000);
                p.kill();
                ret = setMode(dis, name, allOutputs);
            }
            if (lowResFallback && ret != 0) {
                ret = doFbOnly(x, y, allOutputs);
            }
            // Try this stuff again in case --newmode succeeded but --addmode now didn't (NVIDIA)
            if (ret != 0 && !bestDisconMode.modeName.isEmpty()) {
                qDebug() << "Ret" << ret << "- Trying to enable best disconnected screen (2)";
                ret = setMode(bestDisconMode.output, bestDisconMode.modeName, allOutputs);
                bestDisconMode.modeName.clear();
            }
            // Now as fallback, try enabling just one connected output
            if (ret != 0 && !connectedOutputs.isEmpty()) {
                qDebug() << "Ret:" << ret << "- Trying to enable one random connected output";
                QString conn = connectedOutputs.first();
                // Add to output
                p.start("xrandr", QStringList() << "--verbose" << "--addmode" << conn << name);
                p.waitForFinished(2000);
                p.kill();
                ret = setMode(conn, name, allOutputs);
            }
            // Try this stuff again in case --newmode succeeded but --addmode now didn't (NVIDIA)
            if (ret != 0 && !bestConMode.modeName.isEmpty()) {
                qDebug() << "Ret" << ret << "- Trying to enable best connected screen (2)";
                ret = setMode(bestConMode.output, bestConMode.modeName, allOutputs);
                bestConMode.modeName.clear();
            }
        }
    }
    emit loginRequest(lines[0], lines[1], lines[2]);
}

static int doFbOnly(int x, int y, const QStringList &allOutputs)
{
    QStringList setMode(QStringLiteral("--verbose"));
    QProcess p;
    p.setProcessChannelMode(QProcess::ForwardedChannels);
    for (const auto &output : allOutputs) {
        setMode << "--output" << output << "--off";
    }
    setMode << "--fb" << QString("%1x%2").arg(x).arg(y);
    // Set all outputs
    qDebug() << setMode;
    p.start("xrandr", setMode);
    p.waitForFinished(2000);
    int ret = p.exitCode();
    p.kill();
    if (ret != 0)
        return ret;
    p.setProcessChannelMode(QProcess::MergedChannels);
    p.start("xrandr", QStringList("--current"));
    p.waitForFinished(1000);
    p.kill();
    QString xrandrOutput = QString::fromLocal8Bit(p.readAll());
    QRegularExpression re("current (\\d+) x (\\d+),", QRegularExpression::MultilineOption);
    QRegularExpressionMatch m = re.match(xrandrOutput);
    if (!m.hasMatch()) {
        qDebug() << "Could not query current screen size after setting --fb";
    } else {
        if (m.captured(1).toInt() == x && m.captured(2).toInt() == y)
            return 0;
        qDebug() << "Screen size after --fb is not" << x << "x" << y << "but" << m.captured();
    }
    return -1;
}

static int setMode(const QString &virtOut, const QString &name, const QStringList &allOutputs)
{
    QStringList setMode(QStringLiteral("--verbose"));
    QProcess p;
    p.setProcessChannelMode(QProcess::ForwardedChannels);
    for (const auto & output : allOutputs) {
        if (output == virtOut)
            continue; // Skip self
        setMode << "--output" << output << "--off";
    }
    setMode << "--output" << virtOut << "--mode" << name;
    // Set all outputs
    qDebug() << setMode;
    p.start("xrandr", setMode);
    p.waitForFinished(2000);
    int ret = p.exitCode();
    p.kill();
    return ret;
}