#include "ollie_runner.h"

#include <KPluginFactory>
#include <QDBusReply>
#include <QDBusConnection>
#include <QDBusInterface>
#include <QDir>
#include <QThread>
#include <QVariantMap>
#include <QDialog>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QTextEdit>
#include <QPushButton>
#include <QClipboard>
#include <QApplication>
#include <QLabel>

K_PLUGIN_CLASS_WITH_JSON(OllieRunner, "metadata.json")

static const char *SERVICE = "org.ollie.SessionManager";
static const char *OBJ_PATH = "/org/ollie/SessionManager";
static const char *IFACE = "org.ollie.SessionManager";

#if QT_VERSION_MAJOR >= 6
OllieRunner::OllieRunner(QObject *parent, const KPluginMetaData &data)
    : AbstractRunner(parent, data)
#else
OllieRunner::OllieRunner(QObject *parent, const KPluginMetaData &data, const QVariantList &args)
    : AbstractRunner(parent, data, args)
#endif
{
    m_iface = new QDBusInterface(SERVICE, OBJ_PATH, IFACE,
                                 QDBusConnection::sessionBus(), this);
}

void OllieRunner::match(RunnerNS::RunnerContext &context)
{
    const QString query = context.query().trimmed();

    // Only activate on "ollie" prefix
    if (!query.startsWith(QLatin1String("ollie ")) && query != QLatin1String("ollie")) {
        return;
    }

    // Bail if daemon isn't running
    if (!m_iface->isValid()) {
        return;
    }

    QString args = query.mid(6).trimmed(); // everything after "ollie "

    if (args.isEmpty() || args == QLatin1String("list")) {
        // List sessions
        QStringList sessions = listSessions();
        if (sessions.isEmpty()) {
            RunnerNS::QueryMatch match(this);
            match.setId("ollie-no-sessions");
            match.setText("No active sessions");
            match.setIconName("dialog-information");
            match.setRelevance(0.5);
            context.addMatch(match);
        } else {
            for (const QString &entry : sessions) {
                QStringList parts = entry.split('\t');
                if (parts.size() < 4) continue;

                RunnerNS::QueryMatch match(this);
                match.setId("ollie-session-" + parts[0]);
                QString state = parts[1] == "idle" ? "○" : "●";
                match.setText(state + " " + parts[0].left(14) + " [" + parts[3] + "]");
                match.setSubtext(parts[2] + " — " + parts[1]);
                match.setIconName("system-run");
                match.setRelevance(0.7);
                match.setData(parts[0]); // session ID
                context.addMatch(match);
            }
        }
        return;
    }

    if (args.startsWith(QLatin1String("kill "))) {
        QString target = args.mid(5).trimmed();
        RunnerNS::QueryMatch match(this);
        match.setId("ollie-kill-" + target);
        match.setText("Kill session: " + target);
        match.setSubtext("Terminate ollie session");
        match.setIconName("process-stop");
        match.setRelevance(0.8);
        match.setData("kill:" + target);
        context.addMatch(match);
        return;
    }

    if (args.startsWith(QLatin1String("rename "))) {
        // Expected: rename <prefix> <new_name>
        // Prefix-matches against listed sessions so full ID isn't needed
        QString rest = args.mid(7).trimmed();
        int space = rest.indexOf(' ');
        if (space > 0) {
            QString prefix = rest.left(space);
            QString newName = rest.mid(space + 1).trimmed();
            if (!newName.isEmpty()) {
                // Find sessions matching the prefix
                QStringList sessions = listSessions();
                for (const QString &entry : sessions) {
                    QStringList parts = entry.split('\t');
                    if (parts.isEmpty()) continue;
                    QString id = parts[0];
                    if (id.startsWith(prefix) || id == prefix) {
                        RunnerNS::QueryMatch match(this);
                        match.setId("ollie-rename-" + id);
                        match.setText("Rename: " + id + " → " + newName);
                        match.setSubtext("Rename ollie session");
                        match.setIconName("edit-rename");
                        match.setRelevance(0.8);
                        match.setData("rename:" + id + ":" + newName);
                        context.addMatch(match);
                    }
                }
            }
        }
        return;
    }

    // Default: one-shot prompt
    RunnerNS::QueryMatch match(this);
    match.setId("ollie-prompt");
    match.setText("Ask Ollie: " + args);
    match.setSubtext("One-shot prompt — result via notification");
    match.setIconName("system-run");
    match.setRelevance(0.9);
    match.setData("prompt:" + args);
    context.addMatch(match);
}

void OllieRunner::run(const RunnerNS::RunnerContext &context, const RunnerNS::QueryMatch &match)
{
    Q_UNUSED(context);
    QString data = match.data().toString();

    if (data.startsWith("kill:")) {
        QString id = data.mid(5);
        killSession(id);
        return;
    }

    if (data.startsWith("rename:")) {
        // format: rename:<id>:<new_name>
        QString rest = data.mid(7);
        int sep = rest.indexOf(':');
        if (sep > 0) {
            QString id = rest.left(sep);
            QString newName = rest.mid(sep + 1);
            m_iface->call("RenameSession", id, newName);
        }
        return;
    }

    if (data.startsWith("prompt:")) {
        QString prompt = data.mid(7);
        QString result = oneShotRun(prompt);

        // Show popup window with result
        auto *dialog = new QDialog(nullptr);
        dialog->setWindowTitle("Ollie — " + prompt.left(40));
        dialog->setAttribute(Qt::WA_DeleteOnClose);
        dialog->resize(520, 380);

        auto *layout = new QVBoxLayout(dialog);

        auto *textEdit = new QTextEdit(dialog);
        textEdit->setReadOnly(true);
        textEdit->setPlainText(result);
        QFont monoFont("monospace");
        monoFont.setPointSize(QApplication::font().pointSize());
        textEdit->setFont(monoFont);
        layout->addWidget(textEdit);

        auto *buttonLayout = new QHBoxLayout();
        buttonLayout->addStretch();

        auto *copyBtn = new QPushButton("Copy", dialog);
        QObject::connect(copyBtn, &QPushButton::clicked, [result]() {
            QApplication::clipboard()->setText(result);
        });
        buttonLayout->addWidget(copyBtn);

        auto *closeBtn = new QPushButton("Close", dialog);
        QObject::connect(closeBtn, &QPushButton::clicked, dialog, &QDialog::close);
        buttonLayout->addWidget(closeBtn);

        layout->addLayout(buttonLayout);

        dialog->show();
        dialog->raise();
        dialog->activateWindow();
        return;
    }

    // Clicked a session from list — not much to do in KRunner
    // Could open the plasmoid in the future
}

// --- D-Bus helpers ---

QString OllieRunner::oneShotRun(const QString &prompt)
{
    // 1. Create session
    QDBusReply<QString> createReply = m_iface->call(
        "CreateSession", QDir::homePath(), QString(), QString(), "default");
    if (!createReply.isValid()) {
        return "Error: could not create session";
    }
    QString sessionId = createReply.value();

    // 2. Submit prompt
    QDBusReply<bool> submitReply = m_iface->call("Submit", sessionId, prompt);
    if (!submitReply.isValid() || !submitReply.value()) {
        killSession(sessionId);
        return "Error: could not submit prompt";
    }

    // 3. Poll for completion (max 30s)
    QString state = "thinking";
    for (int i = 0; i < 60; i++) {
        QThread::msleep(500);

        QDBusReply<QString> stateReply = m_iface->call("GetState", sessionId);
        if (!stateReply.isValid()) break;
        state = stateReply.value();
        if (state == "idle") break;
    }

    // 4. Get the chat output
    QDBusMessage msg = QDBusMessage::createMethodCall(SERVICE, OBJ_PATH, IFACE, "GetChat");
    msg << sessionId << (qlonglong)0;
    QDBusMessage reply = QDBusConnection::sessionBus().call(msg);

    QString chatText;
    if (reply.type() == QDBusMessage::ReplyMessage && reply.arguments().size() >= 1) {
        chatText = reply.arguments()[0].toString();
    }

    // 5. Extract assistant response (text between [assistant]\n and next \n\n[)
    QString result;
    int aStart = chatText.indexOf("[assistant]\n");
    if (aStart >= 0) {
        aStart += 12; // skip "[assistant]\n"
        int aEnd = chatText.indexOf("\n\n[", aStart);
        if (aEnd < 0) aEnd = chatText.length();
        result = chatText.mid(aStart, aEnd - aStart).trimmed();
    } else {
        result = chatText; // fallback: show everything
    }

    // 6. Kill session
    killSession(sessionId);

    return result.isEmpty() ? "(no response)" : result;
}

QStringList OllieRunner::listSessions()
{
    QDBusReply<QStringList> reply = m_iface->call("ListSessions");
    if (reply.isValid()) return reply.value();
    return {};
}

void OllieRunner::killSession(const QString &id)
{
    m_iface->call("KillSession", id);
}

#include "ollie_runner.moc"