#include "ollie9pclient.h"

#include <QDir>
#include <QStandardPaths>

// Resolve the 9p binary and server address
static QString ninepBin()
{
    // plan9port's 9p command
    QString bin = QStandardPaths::findExecutable("9p");
    if (!bin.isEmpty()) return bin;
    if (QFileInfo::exists("/usr/local/plan9/bin/9p"))
        return "/usr/local/plan9/bin/9p";
    return "9p";
}

static QString serverAddr()
{
    QString ns = qEnvironmentVariable("NAMESPACE");
    if (ns.isEmpty()) {
        // Plan 9 default: /tmp/ns.$USER.$DISPLAY
        QString user = qEnvironmentVariable("USER");
        QString display = qEnvironmentVariable("DISPLAY");
        ns = "/tmp/ns." + user + "." + display;
    }
    return "unix!" + ns + "/ollie";
}

Ollie9pClient::Ollie9pClient(QObject *parent)
    : QObject(parent)
{
    refreshSessions();
}

Ollie9pClient::~Ollie9pClient()
{
    stopStreams();
}

void Ollie9pClient::setActiveSessionId(const QString &id)
{
    if (m_activeSessionId == id) return;

    stopStreams();
    m_activeSessionId = id;
    m_agentId.clear();
    m_activeState = "idle";
    emit activeSessionIdChanged();
    emit activeStateChanged();

    if (id.isEmpty()) return;

    // Discover agent ID
    QByteArray out = run9p({"ls", "session/" + id + "/agent"});
    QString agents = QString::fromUtf8(out).trimmed();
    if (!agents.isEmpty()) {
        m_agentId = agents.split('\n').first().trimmed();
    } else {
        m_agentId = "0";
    }

    startChatStream();
    startStateStream();
}

void Ollie9pClient::refreshSessions()
{
    QByteArray out = run9p({"read", "session/idx"});
    QString raw = QString::fromUtf8(out).trimmed();

    m_sessions.clear();
    if (!raw.isEmpty()) {
        for (const QString &line : raw.split('\n')) {
            QStringList parts = line.split('\t');
            if (parts.size() >= 2) {
                QVariantMap session;
                session["id"] = parts[0];
                session["state"] = parts[1];
                session["model"] = parts.size() > 2 ? parts[2] : "";
                session["agent"] = parts.size() > 3 ? parts[3] : "";
                m_sessions.append(session);
            }
        }
    }
    emit sessionsChanged();

    // Auto-select first session if none active
    if (m_activeSessionId.isEmpty() && !m_sessions.isEmpty()) {
        setActiveSessionId(m_sessions.first().toMap()["id"].toString());
    }
}

QString Ollie9pClient::readLog()
{
    if (m_activeSessionId.isEmpty()) return {};
    QByteArray out = run9p({"read", agentPath() + "/log"});
    return QString::fromUtf8(out);
}

bool Ollie9pClient::submit(const QString &prompt)
{
    if (m_activeSessionId.isEmpty() || prompt.trimmed().isEmpty()) return false;
    QProcess proc;
    proc.start(ninepBin(), {"-a", serverAddr(), "write", agentPath() + "/prompt"});
    proc.waitForStarted(3000);
    proc.write(prompt.toUtf8());
    proc.closeWriteChannel();
    proc.waitForFinished(5000);
    return proc.exitCode() == 0;
}

bool Ollie9pClient::interrupt()
{
    if (m_activeSessionId.isEmpty()) return false;
    QProcess proc;
    proc.start(ninepBin(), {"-a", serverAddr(), "write", agentPath() + "/ctl"});
    proc.waitForStarted(3000);
    proc.write("stop");
    proc.closeWriteChannel();
    proc.waitForFinished(3000);
    return proc.exitCode() == 0;
}

bool Ollie9pClient::kill()
{
    if (m_activeSessionId.isEmpty()) return false;
    QProcess proc;
    proc.start(ninepBin(), {"-a", serverAddr(), "write", agentPath() + "/ctl"});
    proc.waitForStarted(3000);
    proc.write("kill");
    proc.closeWriteChannel();
    proc.waitForFinished(3000);
    return proc.exitCode() == 0;
}

QString Ollie9pClient::getConfig()
{
    if (m_activeSessionId.isEmpty()) return {};
    QByteArray out = run9p({"read", agentPath() + "/cfg"});
    return QString::fromUtf8(out);
}

// --- Streaming ---

void Ollie9pClient::startChatStream()
{
    if (m_chatProc) return;
    m_chatProc = new QProcess(this);
    connect(m_chatProc, &QProcess::readyReadStandardOutput, this, &Ollie9pClient::onChatReadyRead);
    connect(m_chatProc, QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished),
            this, &Ollie9pClient::onChatFinished);
    m_chatProc->start(ninepBin(), {"-a", serverAddr(), "read", agentPath() + "/chat"});
}

void Ollie9pClient::startStateStream()
{
    if (m_stateProc) return;
    m_stateProc = new QProcess(this);
    connect(m_stateProc, &QProcess::readyReadStandardOutput, this, &Ollie9pClient::onStateReadyRead);
    connect(m_stateProc, QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished),
            this, &Ollie9pClient::onStateFinished);
    m_stateProc->start(ninepBin(), {"-a", serverAddr(), "read", agentPath() + "/statewait"});
}

void Ollie9pClient::stopStreams()
{
    if (m_chatProc) {
        m_chatProc->disconnect();
        m_chatProc->kill();
        m_chatProc->waitForFinished(1000);
        m_chatProc->deleteLater();
        m_chatProc = nullptr;
    }
    if (m_stateProc) {
        m_stateProc->disconnect();
        m_stateProc->kill();
        m_stateProc->waitForFinished(1000);
        m_stateProc->deleteLater();
        m_stateProc = nullptr;
    }
}

void Ollie9pClient::onChatReadyRead()
{
    QByteArray data = m_chatProc->readAllStandardOutput();
    if (!data.isEmpty()) {
        emit chatReceived(QString::fromUtf8(data));
    }
}

void Ollie9pClient::onStateReadyRead()
{
    QByteArray data = m_stateProc->readAllStandardOutput();
    QString state = QString::fromUtf8(data).trimmed();
    if (!state.isEmpty()) {
        // statewait returns one line per state change; take the last one
        QStringList lines = state.split('\n', Qt::SkipEmptyParts);
        if (!lines.isEmpty()) {
            m_activeState = lines.last().trimmed();
            emit activeStateChanged();
        }
    }
    // statewait exits after one read — restart it for continuous monitoring
    m_stateProc->deleteLater();
    m_stateProc = nullptr;
    startStateStream();
}

void Ollie9pClient::onChatFinished(int, QProcess::ExitStatus)
{
    // Chat stream ended — session was killed or connection lost
    m_chatProc->deleteLater();
    m_chatProc = nullptr;
    emit streamingDone();
}

void Ollie9pClient::onStateFinished(int, QProcess::ExitStatus)
{
    // State process exited unexpectedly (not via onStateReadyRead restart)
    if (m_stateProc) {
        m_stateProc->deleteLater();
        m_stateProc = nullptr;
    }
}

// --- Helpers ---

QString Ollie9pClient::agentPath() const
{
    return "session/" + m_activeSessionId + "/agent/" + m_agentId;
}

QByteArray Ollie9pClient::run9p(const QStringList &args)
{
    QProcess proc;
    QStringList fullArgs = {"-a", serverAddr()};
    fullArgs.append(args);
    proc.start(ninepBin(), fullArgs);
    proc.waitForFinished(10000);
    return proc.readAllStandardOutput();
}