#include "imapconnection.h"
#include <QDebug>
#include <QEventLoop>
#include <QTimer>

ImapConnection::ImapConnection(QObject *parent) : QObject(parent),
    m_socket(new QSslSocket(this)),
    m_tagCounter(1)
{
    connect(m_socket, &QAbstractSocket::connected, this, &ImapConnection::onConnected);
    connect(m_socket, &QSslSocket::encrypted, this, &ImapConnection::onEncrypted);
    m_readyReadConnection = connect(m_socket, &QSslSocket::readyRead, this, &ImapConnection::onReadyRead);
    connect(m_socket, &QAbstractSocket::errorOccurred, this, &ImapConnection::onSocketError);
}

ImapConnection::~ImapConnection()
{
    disconnect();
}

void ImapConnection::connectToHost(const QString &host, int port, bool useSsl)
{
    m_host = host;
    m_port = port;
    m_useSsl = useSsl;
    if (useSsl) {
        m_socket->connectToHostEncrypted(host, port);
    } else {
        m_socket->connectToHost(host, port);
    }
}

void ImapConnection::login(const QString &username, const QString &password,
                           bool useAuthPlain,
                           std::function<void(bool, const QString&)> callback)
{
    if (!m_socket->isEncrypted() && m_useSsl) {
        callback(false, "Not encrypted");
        return;
    }
    QString cmd;
    if (useAuthPlain) {
        QByteArray auth;
        auth.append('\0');
        auth.append(username.toUtf8());
        auth.append('\0');
        auth.append(password.toUtf8());
        cmd = QString("AUTHENTICATE PLAIN %1").arg(QString::fromLatin1(auth.toBase64()));
    } else {
        cmd = QString("LOGIN %1 %2").arg(username, password);
    }
    sendCommand(cmd, callback);
}

void ImapConnection::sendCommand(const QString &command,
                                std::function<void(bool, const QString&)> callback)
{
    QString tag = generateTag();
    QString fullCmd = tag + " " + command + "\r\n";
    m_socket->write(fullCmd.toUtf8());
    m_socket->flush();
    m_pendingCallbacks[tag] = callback;
}

void ImapConnection::disconnect()
{
    if (m_socket->state() != QAbstractSocket::UnconnectedState) {
        m_socket->disconnectFromHost();
        m_socket->waitForDisconnected(1000);
    }
    m_pendingCallbacks.clear();
}

void ImapConnection::onConnected()
{
    // For non-SSL connections, signal connected immediately.
    // For SSL connections, onEncrypted() fires instead after the SSL handshake.
    if (!m_useSsl) {
        emit connected();
    }
}

void ImapConnection::onEncrypted()
{
    emit connected();
}

void ImapConnection::onReadyRead()
{
    while (m_socket->canReadLine()) {
        QString line = QString::fromUtf8(m_socket->readLine()).trimmed();
        processLine(line);
    }
}

void ImapConnection::onSocketError(QAbstractSocket::SocketError error)
{
    emit errorOccurred(m_socket->errorString());
}

void ImapConnection::processLine(const QString &line)
{
    // If it's an untagged response (starts with *), emit signal
    if (line.startsWith('*')) {
        emit untaggedResponse(line);
        return;
    }

    // Extract the tag prefix (first word) and look it up in pending callbacks
    int spacePos = line.indexOf(' ');
    if (spacePos == -1) return; // malformed line

    QString tagFromLine = line.left(spacePos);

    auto it = m_pendingCallbacks.find(tagFromLine);
    if (it != m_pendingCallbacks.end()) {
        bool ok = line.contains(" OK ");
        auto cb = it.value();
        m_pendingCallbacks.erase(it);
        cb(ok, line);
        return;
    }

    // If we get here, it's an unhandled response
    qWarning() << "Unhandled IMAP response:" << line;
}

void ImapConnection::sendRaw(const QString &data)
{
    m_socket->write(data.toUtf8() + "\r\n");
    m_socket->flush();
}

QString ImapConnection::generateTag()
{
    // Simple tag: A001, A002, ... (like traditional IMAP clients)
    return QString("A%1").arg(m_tagCounter++, 4, 10, QChar('0'));
}

bool ImapConnection::sendCommandWait(const QString &command, QString &response, int msecs)
{
    QByteArray rawResponse;
    const bool completed = sendCommandWaitBytes(command, rawResponse, msecs);
    response = QString::fromUtf8(rawResponse);
    return completed;
}

bool ImapConnection::sendCommandWaitBytes(const QString &command, QByteArray &response, int msecs)
{
    QString tag = generateTag();
    QString fullCmd = tag + " " + command + "\r\n";
    bool complete = false;
    bool timedOut = false;
    QEventLoop loop;
    QTimer timer;
    QByteArray buffer;
    const QByteArray tagBytes = tag.toUtf8();

    // Fetch responses contain arbitrary MIME bytes and IMAP literals. The
    // line-based parser cannot safely consume those bytes, so temporarily
    // take ownership of readyRead and collect the complete tagged response.
    QObject::disconnect(m_readyReadConnection);
    auto collectResponse = [&]() {
        buffer.append(m_socket->readAll());

        int pos = 0;
        while (pos < buffer.size()) {
            const int lineEnd = buffer.indexOf("\r\n", pos);
            if (lineEnd < 0) break;
            const QByteArray line = buffer.mid(pos, lineEnd - pos);
            const int open = line.lastIndexOf('{');
            QByteArray literalSize = open >= 0 && line.endsWith('}')
                ? line.mid(open + 1, line.size() - open - 2) : QByteArray();
            if (literalSize.endsWith('+')) literalSize.chop(1);
            bool literalOk = false;
            const qint64 parsedLiteralSize = literalSize.toLongLong(&literalOk);
            const bool hasLiteral = open >= 0 && line.endsWith('}') && literalOk && parsedLiteralSize >= 0;
            if (hasLiteral) {
                const qint64 length = parsedLiteralSize;
                const int dataStart = lineEnd + 2;
                if (buffer.size() < dataStart + length) break;
                pos = dataStart + int(length);
                if (buffer.mid(pos, 2) == QByteArrayLiteral("\r\n")) pos += 2;
                continue;
            }
            if (line.startsWith(tagBytes + QByteArrayLiteral(" "))) {
                complete = true;
                loop.quit();
                break;
            }
            pos = lineEnd + 2;
        }
    };
    QMetaObject::Connection readyConn = connect(m_socket, &QSslSocket::readyRead, &loop, collectResponse);
    QMetaObject::Connection errorConn = connect(m_socket, &QAbstractSocket::errorOccurred, &loop, [&](QAbstractSocket::SocketError) {
        loop.quit();
    });
    timer.setSingleShot(true);
    QMetaObject::Connection timeoutConn = connect(&timer, &QTimer::timeout, &loop, [&]() {
        timedOut = true;
        loop.quit();
    });

    m_socket->write(fullCmd.toUtf8());
    m_socket->flush();
    if (m_socket->bytesAvailable() > 0) collectResponse();
    timer.start(msecs);
    loop.exec();
    QObject::disconnect(readyConn);
    QObject::disconnect(errorConn);
    QObject::disconnect(timeoutConn);
    m_readyReadConnection = connect(m_socket, &QSslSocket::readyRead, this, &ImapConnection::onReadyRead);

    response = buffer;
    if (timedOut || !complete) return false;
    return true;
}