#include "ssh_session_backend.h"

#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QProcessEnvironment>
#include <QTextStream>
#include <QUuid>

namespace {
QString escapeForShellSingleQuotes(const QString& value)
{
    QString escaped = value;
    escaped.replace(QStringLiteral("'"), QStringLiteral("'\"'\"'"));
    return escaped;
}
}

SshSessionBackend::SshSessionBackend(const Profile& profile, QObject* parent)
    : SessionBackend(profile, parent),
      m_process(new QProcess(this)),
      m_connectedProbeTimer(new QTimer(this)),
      m_state(SessionState::Disconnected),
      m_userInitiatedDisconnect(false),
      m_reconnectPending(false),
      m_waitingForPasswordPrompt(false),
      m_waitingForHostKeyConfirmation(false),
      m_passwordSubmitted(false),
      m_terminalColumns(0),
      m_terminalRows(0)
{
    m_connectedProbeTimer->setSingleShot(true);

    connect(m_process, &QProcess::started, this, &SshSessionBackend::onProcessStarted);
    connect(m_process,
            &QProcess::errorOccurred,
            this,
            &SshSessionBackend::onProcessErrorOccurred);
    connect(m_process,
            qOverload<int, QProcess::ExitStatus>(&QProcess::finished),
            this,
            &SshSessionBackend::onProcessFinished);
    connect(m_process,
            &QProcess::readyReadStandardOutput,
            this,
            &SshSessionBackend::onReadyReadStandardOutput);
    connect(m_process,
            &QProcess::readyReadStandardError,
            this,
            &SshSessionBackend::onReadyReadStandardError);
    connect(m_connectedProbeTimer,
            &QTimer::timeout,
            this,
            &SshSessionBackend::onConnectedProbeTimeout);
}

SshSessionBackend::~SshSessionBackend()
{
    if (m_process->state() != QProcess::NotRunning) {
        m_process->kill();
        m_process->waitForFinished(500);
    }
    cleanupAskPassScript();
}

void SshSessionBackend::connectSession(const SessionConnectOptions& options)
{
    if (m_state == SessionState::Connected || m_state == SessionState::Connecting) {
        emit eventLogged(QStringLiteral("Connect skipped: session is already active."));
        return;
    }

    m_userInitiatedDisconnect = false;
    m_reconnectPending = false;
    m_lastRawError.clear();
    m_activeOptions = options;
    m_waitingForPasswordPrompt = false;
    m_waitingForHostKeyConfirmation = false;
    m_passwordSubmitted = false;

    if (!startSshProcess(options)) {
        return;
    }

    setState(SessionState::Connecting, QStringLiteral("Connecting to SSH endpoint..."));
    emit eventLogged(QStringLiteral("Launching ssh client."));
}

void SshSessionBackend::disconnectSession()
{
    if (m_process->state() == QProcess::NotRunning) {
        if (m_state != SessionState::Disconnected) {
            setState(SessionState::Disconnected, QStringLiteral("Session is disconnected."));
        }
        return;
    }

    m_userInitiatedDisconnect = true;
    emit eventLogged(QStringLiteral("Disconnect requested."));
    m_connectedProbeTimer->stop();

    m_process->terminate();
    QTimer::singleShot(1500,
                       this,
                       [this]() {
                           if (m_process->state() != QProcess::NotRunning) {
                               emit eventLogged(QStringLiteral("Force-stopping ssh process."));
                               m_process->kill();
                           }
                       });
}

void SshSessionBackend::reconnectSession(const SessionConnectOptions& options)
{
    emit eventLogged(QStringLiteral("Reconnect requested."));

    if (m_process->state() == QProcess::NotRunning) {
        connectSession(options);
        return;
    }

    m_reconnectPending = true;
    m_reconnectOptions = options;
    m_userInitiatedDisconnect = true;
    m_process->terminate();
}

void SshSessionBackend::sendInput(const QString& input)
{
    if (m_process->state() != QProcess::Running) {
        emit eventLogged(QStringLiteral("Input ignored: session is not running."));
        return;
    }

    if (input.isEmpty()) {
        return;
    }

    m_process->write(input.toUtf8());
}

void SshSessionBackend::confirmHostKey(bool trustHost)
{
    if (m_process->state() != QProcess::Running || !m_waitingForHostKeyConfirmation) {
        return;
    }

    m_waitingForHostKeyConfirmation = false;
    const QString response = trustHost ? QStringLiteral("yes\n") : QStringLiteral("no\n");
    m_process->write(response.toUtf8());

    emit eventLogged(trustHost
                         ? QStringLiteral("Host key accepted by user.")
                         : QStringLiteral("Host key rejected by user."));
}

void SshSessionBackend::updateTerminalSize(int columns, int rows)
{
    m_terminalColumns = columns;
    m_terminalRows = rows;

    if (m_state == SessionState::Connected) {
        applyTerminalSizeIfAvailable();
    }
}

void SshSessionBackend::onProcessStarted()
{
    emit eventLogged(QStringLiteral("ssh process started."));
    m_connectedProbeTimer->start(1200);
}

void SshSessionBackend::onProcessErrorOccurred(QProcess::ProcessError)
{
    const QString rawError = m_process->errorString();
    if (!rawError.isEmpty()) {
        m_lastRawError += rawError + QLatin1Char('\n');
    }

    if (m_state == SessionState::Connecting) {
        const QString display = mapSshError(m_lastRawError);
        setState(SessionState::Failed, display);
        emit connectionError(display, m_lastRawError.trimmed());
    }
}

void SshSessionBackend::onProcessFinished(int exitCode, QProcess::ExitStatus)
{
    m_connectedProbeTimer->stop();
    cleanupAskPassScript();

    if (m_reconnectPending) {
        m_reconnectPending = false;
        SessionConnectOptions options = m_reconnectOptions;
        setState(SessionState::Disconnected, QStringLiteral("Reconnecting..."));
        QTimer::singleShot(0, this, [this, options]() { connectSession(options); });
        return;
    }

    if (m_userInitiatedDisconnect) {
        m_userInitiatedDisconnect = false;
        setState(SessionState::Disconnected, QStringLiteral("Session disconnected."));
        emit eventLogged(QStringLiteral("ssh process exited after disconnect request."));
        return;
    }

    if (m_state == SessionState::Connecting || exitCode != 0) {
        QString rawError = m_lastRawError.trimmed();
        if (rawError.isEmpty()) {
            rawError = QStringLiteral("ssh exited with code %1").arg(exitCode);
        }

        const QString display = mapSshError(rawError);
        setState(SessionState::Failed, display);
        emit connectionError(display, rawError);
        return;
    }

    setState(SessionState::Disconnected, QStringLiteral("SSH session ended."));
}

void SshSessionBackend::onReadyReadStandardOutput()
{
    const QString chunk = QString::fromUtf8(m_process->readAllStandardOutput());
    if (chunk.isEmpty()) {
        return;
    }

    emit outputReceived(chunk);

    if (m_state == SessionState::Connecting && !m_waitingForHostKeyConfirmation
        && !m_waitingForPasswordPrompt) {
        setState(SessionState::Connected, QStringLiteral("SSH session established."));
    }
}

void SshSessionBackend::onReadyReadStandardError()
{
    const QString chunk = QString::fromUtf8(m_process->readAllStandardError());
    if (chunk.isEmpty()) {
        return;
    }

    m_lastRawError += chunk;
    emit outputReceived(chunk);

    const QStringList lines = chunk.split(QLatin1Char('\n'), Qt::SkipEmptyParts);
    for (const QString& line : lines) {
        const QString trimmed = line.trimmed();
        if (!trimmed.isEmpty()) {
            emit eventLogged(trimmed);
        }

        if (trimmed.contains(QStringLiteral("Are you sure you want to continue connecting"),
                             Qt::CaseInsensitive)
            && !m_waitingForHostKeyConfirmation) {
            m_waitingForHostKeyConfirmation = true;
            emit eventLogged(QStringLiteral("Awaiting host key confirmation from user."));
            emit hostKeyConfirmationRequested(trimmed);
            continue;
        }

        if (trimmed.contains(QStringLiteral("password:"), Qt::CaseInsensitive)
            && profile().authMode.compare(QStringLiteral("Password"), Qt::CaseInsensitive) == 0
            && !m_passwordSubmitted) {
            if (m_activeOptions.password.isEmpty()) {
                const QString message = QStringLiteral("Password prompt received but no password is available.");
                setState(SessionState::Failed, message);
                emit connectionError(message, trimmed);
                return;
            }

            m_waitingForPasswordPrompt = false;
            m_passwordSubmitted = true;
            m_process->write((m_activeOptions.password + QStringLiteral("\n")).toUtf8());
            emit eventLogged(QStringLiteral("Password prompt received; credentials submitted."));
            continue;
        }
    }
}

void SshSessionBackend::onConnectedProbeTimeout()
{
    if (m_state != SessionState::Connecting) {
        return;
    }

    if (m_process->state() == QProcess::Running && !m_waitingForHostKeyConfirmation
        && !m_waitingForPasswordPrompt) {
        setState(SessionState::Connected, QStringLiteral("SSH session established."));
    }
}

void SshSessionBackend::setState(SessionState state, const QString& message)
{
    m_state = state;
    emit stateChanged(state, message);
    emit eventLogged(message);

    if (m_state == SessionState::Connected) {
        applyTerminalSizeIfAvailable();
    }
}

bool SshSessionBackend::startSshProcess(const SessionConnectOptions& options)
{
    const Profile& p = profile();

    if (p.host.trimmed().isEmpty()) {
        const QString message = QStringLiteral("Host is required for SSH connections.");
        setState(SessionState::Failed, message);
        emit connectionError(message, message);
        return false;
    }

    if (p.port < 1 || p.port > 65535) {
        const QString message = QStringLiteral("Port must be between 1 and 65535.");
        setState(SessionState::Failed, message);
        emit connectionError(message, message);
        return false;
    }

    QStringList args;
    args << QStringLiteral("-tt") << QStringLiteral("-p") << QString::number(p.port)
         << QStringLiteral("-o") << QStringLiteral("ConnectTimeout=12") << QStringLiteral("-o")
         << QStringLiteral("ServerAliveInterval=20") << QStringLiteral("-o")
         << QStringLiteral("ServerAliveCountMax=2");

    const QString policy = options.knownHostsPolicy.trimmed().isEmpty()
        ? p.knownHostsPolicy.trimmed()
        : options.knownHostsPolicy.trimmed();

    if (policy.compare(QStringLiteral("Ignore"), Qt::CaseInsensitive) == 0) {
        args << QStringLiteral("-o") << QStringLiteral("StrictHostKeyChecking=no")
             << QStringLiteral("-o")
             << QStringLiteral("UserKnownHostsFile=%1").arg(knownHostsFileForNullDevice());
    } else if (policy.compare(QStringLiteral("Accept New"), Qt::CaseInsensitive) == 0) {
        args << QStringLiteral("-o") << QStringLiteral("StrictHostKeyChecking=accept-new");
    } else if (policy.compare(QStringLiteral("Ask"), Qt::CaseInsensitive) == 0) {
        args << QStringLiteral("-o") << QStringLiteral("StrictHostKeyChecking=ask");
    } else {
        args << QStringLiteral("-o") << QStringLiteral("StrictHostKeyChecking=yes");
    }

    QProcessEnvironment environment = QProcessEnvironment::systemEnvironment();

    if (p.authMode.compare(QStringLiteral("Password"), Qt::CaseInsensitive) == 0) {
        if (options.password.isEmpty()) {
            const QString message = QStringLiteral("Password is required for password authentication.");
            setState(SessionState::Failed, message);
            emit connectionError(message, message);
            return false;
        }

        args << QStringLiteral("-o") << QStringLiteral("PreferredAuthentications=password")
             << QStringLiteral("-o") << QStringLiteral("PubkeyAuthentication=no")
             << QStringLiteral("-o") << QStringLiteral("NumberOfPasswordPrompts=1");
        m_waitingForPasswordPrompt = false;

        QString askPassError;
        if (!configureAskPass(options, environment, askPassError)) {
            setState(SessionState::Failed, askPassError);
            emit connectionError(askPassError, askPassError);
            return false;
        }
    } else if (p.authMode.compare(QStringLiteral("Private Key"), Qt::CaseInsensitive) == 0) {
        QString keyPath = options.privateKeyPath.trimmed();
        if (keyPath.isEmpty()) {
            keyPath = p.privateKeyPath.trimmed();
        }

        if (keyPath.isEmpty()) {
            const QString message = QStringLiteral("Private key path is required.");
            setState(SessionState::Failed, message);
            emit connectionError(message, message);
            return false;
        }

        if (!QFileInfo::exists(keyPath)) {
            const QString message = QStringLiteral("Private key file does not exist: %1")
                                        .arg(keyPath);
            setState(SessionState::Failed, message);
            emit connectionError(message, message);
            return false;
        }

        args << QStringLiteral("-i") << keyPath << QStringLiteral("-o")
             << QStringLiteral("PreferredAuthentications=publickey") << QStringLiteral("-o")
             << QStringLiteral("PasswordAuthentication=no");
    }

    const QString target = p.username.trimmed().isEmpty()
        ? p.host.trimmed()
        : QStringLiteral("%1@%2").arg(p.username.trimmed(), p.host.trimmed());
    args << target;

    m_process->setProcessEnvironment(environment);
    m_process->setProgram(QStringLiteral("ssh"));
    m_process->setArguments(args);
    m_process->setProcessChannelMode(QProcess::SeparateChannels);

    m_process->start();
    if (!m_process->waitForStarted(3000)) {
        const QString rawError = m_process->errorString();
        const QString display = mapSshError(rawError);
        setState(SessionState::Failed, display);
        emit connectionError(display, rawError);
        return false;
    }

    return true;
}

bool SshSessionBackend::configureAskPass(const SessionConnectOptions& options,
                                         QProcessEnvironment& environment,
                                         QString& error)
{
    cleanupAskPassScript();

#ifdef Q_OS_WIN
    m_askPassScriptPath = QDir::temp().filePath(
        QStringLiteral("orbithub_askpass_%1.cmd")
            .arg(QUuid::createUuid().toString(QUuid::WithoutBraces)));
#else
    m_askPassScriptPath = QDir::temp().filePath(
        QStringLiteral("orbithub_askpass_%1.sh")
            .arg(QUuid::createUuid().toString(QUuid::WithoutBraces)));
#endif

    QFile script(m_askPassScriptPath);
    if (!script.open(QIODevice::WriteOnly | QIODevice::Truncate | QIODevice::Text)) {
        error = QStringLiteral("Failed to create temporary askpass helper script.");
        cleanupAskPassScript();
        return false;
    }

    QTextStream out(&script);
#ifdef Q_OS_WIN
    out << "@echo off\r\n";
    out << "echo " << options.password << "\r\n";
#else
    const QString escapedPassword = escapeForShellSingleQuotes(options.password);
    out << "#!/bin/sh\n";
    out << "printf '%s\\n' '" << escapedPassword << "'\n";
#endif
    out.flush();
    script.close();

#ifndef Q_OS_WIN
    if (!QFile::setPermissions(m_askPassScriptPath,
                               QFile::ReadOwner | QFile::WriteOwner | QFile::ExeOwner)) {
        error = QStringLiteral("Failed to set permissions on askpass helper script.");
        cleanupAskPassScript();
        return false;
    }
#endif

    environment.insert(QStringLiteral("SSH_ASKPASS"), m_askPassScriptPath);
    environment.insert(QStringLiteral("SSH_ASKPASS_REQUIRE"), QStringLiteral("force"));
    if (!environment.contains(QStringLiteral("DISPLAY"))) {
        environment.insert(QStringLiteral("DISPLAY"), QStringLiteral(":0"));
    }
    return true;
}

void SshSessionBackend::cleanupAskPassScript()
{
    if (!m_askPassScriptPath.isEmpty()) {
        QFile::remove(m_askPassScriptPath);
        m_askPassScriptPath.clear();
    }
}

QString SshSessionBackend::mapSshError(const QString& rawError) const
{
    const QString raw = rawError.trimmed();
    if (raw.contains(QStringLiteral("Permission denied"), Qt::CaseInsensitive)) {
        return QStringLiteral("Authentication failed. Check username and credentials.");
    }
    if (raw.contains(QStringLiteral("Host key verification failed"), Qt::CaseInsensitive)) {
        return QStringLiteral("Host key verification failed.");
    }
    if (raw.contains(QStringLiteral("Could not resolve hostname"), Qt::CaseInsensitive)) {
        return QStringLiteral("Host could not be resolved.");
    }
    if (raw.contains(QStringLiteral("Connection timed out"), Qt::CaseInsensitive)
        || raw.contains(QStringLiteral("Operation timed out"), Qt::CaseInsensitive)) {
        return QStringLiteral("Connection timed out.");
    }
    if (raw.contains(QStringLiteral("Connection refused"), Qt::CaseInsensitive)) {
        return QStringLiteral("Connection refused by remote host.");
    }
    if (raw.contains(QStringLiteral("No route to host"), Qt::CaseInsensitive)) {
        return QStringLiteral("No route to host.");
    }
    if (raw.contains(QStringLiteral("Identity file"), Qt::CaseInsensitive)
        && raw.contains(QStringLiteral("not accessible"), Qt::CaseInsensitive)) {
        return QStringLiteral("Private key file is not accessible.");
    }
    if (raw.contains(QStringLiteral("No such file or directory"), Qt::CaseInsensitive)) {
        if (raw.contains(QStringLiteral("ssh-askpass"), Qt::CaseInsensitive)) {
            return QStringLiteral("SSH password helper is missing or failed to launch.");
        }
        return QStringLiteral("Required file was not found.");
    }
    if (raw.isEmpty()) {
        return QStringLiteral("SSH connection failed for an unknown reason.");
    }

    return QStringLiteral("SSH connection failed.");
}

QString SshSessionBackend::knownHostsFileForNullDevice() const
{
#ifdef Q_OS_WIN
    return QStringLiteral("NUL");
#else
    return QStringLiteral("/dev/null");
#endif
}

void SshSessionBackend::applyTerminalSizeIfAvailable()
{
    if (m_process->state() != QProcess::Running) {
        return;
    }

    if (m_terminalColumns <= 0 || m_terminalRows <= 0) {
        return;
    }

    const QString command = QStringLiteral("stty cols %1 rows %2\\n")
                                .arg(m_terminalColumns)
                                .arg(m_terminalRows);
    m_process->write(command.toUtf8());
    emit eventLogged(
        QStringLiteral("Applied terminal size: %1x%2").arg(m_terminalColumns).arg(m_terminalRows));
}