z-bar-qt/Plugins/ZShell/Llm/session.cpp

#include "session.hpp"

#include "attachment.hpp"
#include "chatstore.hpp"
#include "llmclient.hpp"

#include <QClipboard>
#include <QDebug>
#include <QDateTime>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QGuiApplication>
#include <QImage>
#include <QMimeDatabase>
#include <QMimeData>
#include <QMimeType>
#include <QStandardPaths>
#include <QUrl>
#include <QtGlobal>

namespace ZShell::llm {

namespace {

constexpr qint64 kMaxImageBytes = 20 * 1024 * 1024;
constexpr qint64 kMaxFileBytes = 100 * 1024 * 1024;

QString titleFrom(const QString& content) {
	const QString flat = content.simplified();
	if (flat.isEmpty()) return QString();
	if (flat.size() <= 48) return flat;
	return flat.left(47) + QStringLiteral("…");
}

} // namespace

ChatSession::ChatSession(const QString& id, QObject* parent)
	: QObject(parent), m_id(id), m_workspace(sandboxDirFor(id)) {
	m_model = new ChatMessageModel(this, this);
	connect(
		m_model,
		&QAbstractListModel::rowsInserted,
		this,
		&ChatSession::onModelRowsChanged);
	connect(
		m_model,
		&QAbstractListModel::rowsRemoved,
		this,
		&ChatSession::onModelRowsChanged);
	connect(
		m_model,
		&QAbstractListModel::modelReset,
		this,
		&ChatSession::onModelRowsChanged);
}

void ChatSession::onModelRowsChanged() {
	setCount(m_model->rowCount());
}

void ChatSession::setTitle(const QString& value) {
	if (m_title == value) return;
	m_title = value;
	Q_EMIT titleChanged();
	persist();
}

void ChatSession::setIcon(const QString& value) {
	if (m_icon == value) return;
	m_icon = value;
	Q_EMIT iconChanged();
}

void ChatSession::setUpdatedAt(qint64 value) {
	if (m_updatedAt == value) return;
	m_updatedAt = value;
	Q_EMIT updatedAtChanged();
}

void ChatSession::setPinned(bool value) {
	if (m_pinned == value) return;
	m_pinned = value;
	Q_EMIT pinnedChanged();
}

void ChatSession::setCount(int value) {
	if (m_messageCount == value) return;
	m_messageCount = value;
	Q_EMIT messageCountChanged();
}

void ChatSession::setMeta(
	const QString& title, qint64 createdAt, qint64 updatedAt, int messageCount) {
	m_title = title;
	m_createdAt = createdAt;
	m_updatedAt = updatedAt;
	m_messageCount = messageCount;
}

void ChatSession::setLastTokenCount(int value) {
	m_lastTokenCount = value;
}

LlmClient* ChatSession::client() const {
	if (auto* store = qobject_cast<ChatStore*>(parent()))
		return store->llmClient();
	return nullptr;
}

void ChatSession::ensureLoaded() {
	if (m_loadRequested) return;
	m_loadRequested = true;
	if (auto* store = qobject_cast<ChatStore*>(parent()))
		store->loadMessagesInto(this);
}

ChatMessageModel* ChatSession::messagesModel() {
	ensureLoaded();
	return m_model;
}

void ChatSession::markLoaded() {
	if (m_loaded) return;
	m_loaded = true;
	Q_EMIT loaded();
}

bool ChatSession::takeClearPending() {
	const bool pending = m_clearPending;
	m_clearPending = false;
	return pending;
}

void ChatSession::adoptMessages(QList<ChatMessage*> messages) {
	m_model->loadMessages(std::move(messages));
}

void ChatSession::persist() {
	if (auto* store = qobject_cast<ChatStore*>(parent())) store->persist(this);
}

ChatMessage* ChatSession::appendNewest(
	ChatMessage::Role role, const QString& content, qint64 timestamp) {
	return m_model->appendNewest(role, content, timestamp);
}

void ChatSession::removeMessage(ChatMessage* message) {
	m_model->removeMessage(message);
}

void ChatSession::clearMessages() {
	m_model->clear();
}

void ChatSession::startGeneration(ChatMessage* target) {
	auto* generation = target->activeGeneration();
	auto* clientObject = client();
	if (generation && clientObject) {
		if (isLoaded()) {
			clientObject->startGeneration(this, generation);
		} else {
			connect(
				this,
				&ChatSession::loaded,
				clientObject,
				[this, generation, clientObject]() {
					clientObject->startGeneration(this, generation);
				});
		}
	}
	persist();
}

QString ChatSession::sandboxDirFor(const QString& sessionId) {
	const QString state =
		QStandardPaths::writableLocation(QStandardPaths::GenericStateLocation);
	return state + QStringLiteral("/zshell/sandboxes/") + sessionId;
}

void ChatSession::ensureWorkspace() {
	QDir().mkpath(m_workspace);
}

void ChatSession::setAttachmentError(const QString& value) {
	if (m_attachmentError == value) return;
	m_attachmentError = value;
	Q_EMIT attachmentErrorChanged();
}

QVariantList ChatSession::addAttachments(const QStringList& paths) {
	ensureWorkspace();
	setAttachmentError(QString());
	QVariantList result;
	for (const QString& rawPath : paths) {
		const QFileInfo info(rawPath);
		if (!info.exists() || !info.isFile()) {
			qWarning() << "ChatSession: addAttachments skipping" << rawPath
					   << ": file not found or not a file";
			setAttachmentError(
				QStringLiteral("File not found: %1").arg(info.fileName()));
			continue;
		}
		const QString mimeType = QMimeDatabase().mimeTypeForFile(info).name();
		const bool isImage =
			mimeType.startsWith(QLatin1String("image/"), Qt::CaseInsensitive);
		const qint64 cap = isImage ? kMaxImageBytes : kMaxFileBytes;
		if (info.size() > cap) {
			qWarning() << "ChatSession: addAttachments skipping" << rawPath
					   << ": too large (" << info.size() << "bytes, cap" << cap
					   << ")";
			setAttachmentError(
				QStringLiteral("File too large: %1").arg(info.fileName()));
			continue;
		}
		QString destination = m_workspace + QLatin1Char('/') + info.fileName();
		int suffix = 1;
		while (QFile::exists(destination)) {
			const QString extension = info.suffix().isEmpty()
										  ? QString()
										  : QLatin1Char('.') + info.suffix();
			destination = m_workspace + QLatin1Char('/') +
						  info.completeBaseName() +
						  QStringLiteral("_%1%2").arg(suffix).arg(extension);
			++suffix;
		}
		if (!QFile::copy(info.absoluteFilePath(), destination)) {
			qWarning() << "ChatSession: addAttachments skipping" << rawPath
					   << ": copy to" << destination << "failed";
			setAttachmentError(
				QStringLiteral("Could not copy file: %1").arg(info.fileName()));
			continue;
		}
		auto* attachment = new Attachment(
			info.fileName(),
			destination,
			isImage ? Attachment::Kind::Image : Attachment::Kind::File,
			mimeType,
			info.size(),
			this);
		result.append(QVariant::fromValue(attachment));
	}
	return result;
}

bool ChatSession::clipboardHasAttachment() {
	QClipboard* clipboard = QGuiApplication::clipboard();
	if (!clipboard) return false;
	const QMimeData* data = clipboard->mimeData();
	if (!data) return false;
	if (data->hasImage()) return true;
	for (const QUrl& url : data->urls())
		if (url.isLocalFile()) return true;
	return false;
}

QVariantList ChatSession::pasteAttachments() {
	QClipboard* clipboard = QGuiApplication::clipboard();
	const QMimeData* data = clipboard ? clipboard->mimeData() : nullptr;
	if (!data) {
		qWarning() << "ChatSession: pasteAttachments: no clipboard data";
		return {};
	}

	qInfo() << "ChatSession: pasteAttachments: formats =" << data->formats()
			<< "hasImage =" << data->hasImage()
			<< "hasUrls =" << data->hasUrls() << "hasText =" << data->hasText();

	QStringList paths;
	QString tempImagePath;
	if (data->hasImage()) {
		const QImage image = clipboard->image();
		qInfo() << "ChatSession: pasteAttachments: image"
				<< (image.isNull() ? "null"
								   : QString::number(image.width()) + "x" +
										 QString::number(image.height()));
		if (image.isNull()) {
			setAttachmentError(
				QStringLiteral("Clipboard image could not be decoded"));
			return {};
		}
		tempImagePath =
			QDir::tempPath() + QStringLiteral("/zshell-pasted-image-%1.png")
								   .arg(QDateTime::currentMSecsSinceEpoch());
		if (!image.save(tempImagePath)) {
			qWarning() << "ChatSession: pasteAttachments: saving image to"
					   << tempImagePath << "failed";
			setAttachmentError(
				QStringLiteral("Could not save the pasted image"));
			return {};
		}
		paths.append(tempImagePath);
	} else if (data->hasUrls()) {
		for (const QUrl& url : data->urls()) {
			qInfo() << "ChatSession: pasteAttachments: url =" << url.toString()
					<< "local =" << url.isLocalFile();
			if (url.isLocalFile()) paths.append(url.toLocalFile());
		}
	}
	if (paths.isEmpty()) {
		if (data->hasImage() || data->hasUrls())
			setAttachmentError(QStringLiteral(
				"Nothing usable on the clipboard (expected an image "
				"or copied files)"));
		return {};
	}
	const QVariantList attachments = addAttachments(paths);
	if (!tempImagePath.isEmpty()) QFile::remove(tempImagePath);
	qInfo() << "ChatSession: pasteAttachments: result =" << attachments.size()
			<< "attachment(s)";
	return attachments;
}

void ChatSession::sendMessage(const QString& text) {
	sendMessage(text, QVariantList());
}

void ChatSession::sendMessage(
	const QString& text, const QVariantList& attachments) {
	QList<Attachment*> attachmentList;
	for (const QVariant& variant : attachments)
		if (auto* attachment =
				qobject_cast<Attachment*>(variant.value<QObject*>()))
			attachmentList.append(attachment);
	const QString trimmed = text.trimmed();
	if (trimmed.isEmpty() && attachmentList.isEmpty()) return;
	if (auto* clientObject = client()) {
		if (clientObject->busy()) return;
	}

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	auto* userMessage = appendNewest(ChatMessage::Role::User, trimmed, now);
	if (!attachmentList.isEmpty()) userMessage->setAttachments(attachmentList);
	while (m_model->rowCount() > 200)
		m_model->removeMessage(m_model->at(m_model->rowCount() - 1));
	if (m_title.isEmpty()) {
		QString titleSource = trimmed;
		if (titleSource.isEmpty() && !attachmentList.isEmpty())
			titleSource = attachmentList.first()->name();
		m_title = titleFrom(titleSource);
		Q_EMIT titleChanged();
		qInfo() << "ChatSession:" << m_id << "new conversation,"
				<< "fallback title" << m_title
				<< "- requesting generated title and icon";
		if (auto* clientObject = client()) {
			clientObject->requestTitle(this, titleSource);
			clientObject->requestIcon(this, titleSource);
		}
	}
	auto* assistant =
		appendNewest(ChatMessage::Role::Assistant, QString(), now);
	startGeneration(assistant);
}

void ChatSession::retry(ChatMessage* target) {
	if (!target || target->role() != ChatMessage::Role::Assistant) return;
	if (auto* clientObject = client()) {
		if (clientObject->busy()) return;
	}
	const int row = m_model->rowOf(target);
	if (row < 0) return;

	m_model->removeRange(0, row - 1);
	if (m_model->rowCount() < 2) return;

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	target->appendGeneration(now);
	startGeneration(target);
}

void ChatSession::continueFrom(ChatMessage* message) {
	if (!message) return;
	if (auto* clientObject = client()) {
		if (clientObject->busy()) return;
	}
	const int row = m_model->rowOf(message);
	if (row < 0) return;

	m_model->removeRange(0, row - 1);

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	if (message->role() == ChatMessage::Role::Assistant) {
		if (m_model->rowCount() < 2) return;
		message->appendGeneration(now);
		startGeneration(message);
	} else {
		auto* assistant =
			appendNewest(ChatMessage::Role::Assistant, QString(), now);
		startGeneration(assistant);
	}
}

void ChatSession::clear() {
	if (auto* clientObject = client()) {
		if (clientObject->busy() && clientObject->streamingSession() == this) {
			clientObject->clearOnFinish(this);
			clientObject->stop();
			return;
		}
	}
	if (!isLoaded()) {
		m_clearPending = true;
		return;
	}
	m_model->clear();
	persist();
}

} // namespace ZShell::llm