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

#include "chatstore.hpp"

#include "chat.hpp"

#include <QDateTime>
#include <QDebug>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QFileInfoList>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QSaveFile>
#include <QStandardPaths>

#include <algorithm>

namespace ZShell {

ChatStore::ChatStore(QObject* parent) : QObject(parent) {
	load();
}

QString ChatStore::dir() {
	const QString base =
		QStandardPaths::writableLocation(QStandardPaths::GenericDataLocation) +
		QStringLiteral("/zshell");
	return base + QStringLiteral("/chats");
}

QString ChatStore::pathFor(const QString& id) {
	return dir() + QStringLiteral("/") + id + QStringLiteral(".json");
}

bool ChatStore::writeFile(const QString& path, const QJsonObject& doc) {
	QDir().mkpath(QFileInfo(path).absolutePath());

	QSaveFile file(path);
	if (!file.open(QIODevice::WriteOnly)) {
		qWarning() << "ChatStore: failed to open" << path << "for writing:"
		           << file.errorString();
		return false;
	}
	if (file.write(QJsonDocument(doc).toJson(QJsonDocument::Indented)) < 0 ||
		!file.commit()) {
		qWarning() << "ChatStore: failed to write" << path << ":"
		           << file.errorString();
		return false;
	}
	return true;
}

int ChatStore::count() const {
	return static_cast<int>(m_sessions.size());
}

QVariantList ChatStore::values() const {
	QVariantList vals;
	vals.reserve(m_sessions.size());
	for (const auto* session : m_sessions)
		vals.append(QVariant::fromValue(session));
	return vals;
}

ChatSession* ChatStore::at(int index) const {
	if (index < 0 || index >= m_sessions.size())
		return nullptr;
	return m_sessions.at(index);
}

ChatSession* ChatStore::insert(int index) {
	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	auto* session = new ChatSession(QString::number(now), this);
	session->setPath(pathFor(session->id()));
	session->setMeta(QString(), now, now, 0);
	const int pos =
		index >= 0 && index <= m_sessions.size() ? index : 0;
	m_sessions.insert(pos, session);
	Q_EMIT countChanged();
	Q_EMIT valuesChanged();
	return session;
}

void ChatStore::remove(int index) {
	removeSession(at(index));
}

void ChatStore::remove(ChatSession* chat) {
	removeSession(chat);
}

void ChatStore::removeSession(ChatSession* session) {
	if (!session || !m_sessions.contains(session))
		return;
	if (isStreaming(session)) {
		qWarning() << "ChatStore: cannot remove a chat that is streaming";
		return;
	}
	const QList<ChatSession*> before = m_sessions;
	QFile::remove(session->path());
	if (auto* chat = qobject_cast<Chat*>(parent()))
		if (chat->m_active == session)
			chat->m_active = nullptr;
	m_sessions.removeOne(session);
	session->deleteLater();
	notify(before);
}

void ChatStore::move(int from, int to) {
	if (from < 0 || from >= m_sessions.size() || to < 0 ||
		to >= m_sessions.size() || from == to)
		return;
	m_sessions.move(from, to);
	Q_EMIT valuesChanged();
}

void ChatStore::clear() {
	const QList<ChatSession*> before = m_sessions;
	QList<ChatSession*> keep;
	for (ChatSession* session : std::as_const(m_sessions)) {
		if (isStreaming(session)) {
			keep.append(session);
			continue;
		}
		QFile::remove(session->path());
		session->deleteLater();
	}
	m_sessions = keep;
	notify(before);
}

ChatSession* ChatStore::sessionById(const QString& id) {
	for (auto* session : m_sessions)
		if (session->id() == id)
			return session;
	return nullptr;
}

void ChatStore::persist(ChatSession* session) {
	if (!session || !m_sessions.contains(session))
		return;
	session->setUpdatedAt(QDateTime::currentMSecsSinceEpoch());
	if (!writeFile(session->path(), session->document()))
		return;
	sortAndNotify();
}

bool ChatStore::isStreaming(ChatSession* session) const {
	const auto* chat = qobject_cast<const Chat*>(parent());
	return chat && chat->streamingSession() == session;
}

void ChatStore::sortAndNotify() {
	const QList<ChatSession*> before = m_sessions;
	std::stable_sort(
		m_sessions.begin(),
		m_sessions.end(),
		[](const ChatSession* a, const ChatSession* b) {
			return a->updatedAtMs() > b->updatedAtMs();
		});
	notify(before);
}

void ChatStore::notify(const QList<ChatSession*>& before) {
	if (before.size() != m_sessions.size())
		Q_EMIT countChanged();
	bool same = before.size() == m_sessions.size();
	for (int i = 0; same && i < m_sessions.size(); ++i)
		if (before.at(i) != m_sessions.at(i))
			same = false;
	if (!same)
		Q_EMIT valuesChanged();
}

void ChatStore::load() {
	const QFileInfoList files = QDir(dir()).entryInfoList(
		QStringList{QStringLiteral("*.json")},
		QDir::Files,
		QDir::Time);
	for (const QFileInfo& info : files) {
		QFile file(info.absoluteFilePath());
		if (!file.open(QIODevice::ReadOnly))
			continue;
		const QJsonDocument doc = QJsonDocument::fromJson(file.readAll());
		file.close();
		if (!doc.isObject())
			continue;
		const QJsonObject obj = doc.object();
		auto* session =
			new ChatSession(obj["id"].toString(info.completeBaseName()), this);
		session->setPath(info.absoluteFilePath());
		session->setMeta(
			obj["title"].toString(),
			obj["createdAt"].toVariant().toLongLong(),
			obj["updatedAt"].toVariant().toLongLong(),
			int(obj["messages"].toArray().size()));
		m_sessions.append(session);
	}
	sortAndNotify();
	migrateLegacy();
}

void ChatStore::migrateLegacy() {
	const QString path = dir() + QStringLiteral("/chat.json");
	QFile file(path);
	if (!file.open(QIODevice::ReadOnly))
		return;
	const QJsonDocument doc = QJsonDocument::fromJson(file.readAll());
	file.close();
	const QJsonArray arr = doc.isArray() ? doc.array() : QJsonArray();

	QString title;
	qint64 first = 0;
	qint64 last = 0;
	QJsonArray messages;
	for (const QJsonValue& value : arr) {
		const QJsonObject msg = value.toObject();
		const QString content = msg["content"].toString();
		if (content.isEmpty())
			continue;
		if (first == 0)
			first = msg["timestamp"].toVariant().toLongLong();
		last = msg["timestamp"].toVariant().toLongLong();
		if (title.isEmpty() && msg["role"].toString() == "user")
			title = Chat::titleFrom(content);
		messages.append(msg);
	}
	if (messages.isEmpty())
		return;

	const QString id = QString::number(QDateTime::currentMSecsSinceEpoch());
	auto* session = new ChatSession(id, this);
	session->setPath(pathFor(id));

	QJsonObject out;
	out[QStringLiteral("id")] = id;
	out[QStringLiteral("title")] = title;
	out[QStringLiteral("createdAt")] = first;
	out[QStringLiteral("updatedAt")] = last;
	out[QStringLiteral("messages")] = messages;
	if (!writeFile(session->path(), out))
		return;
	QFile::remove(path);

	session->setMeta(title, first, last, int(messages.size()));
	m_sessions.append(session);
	sortAndNotify();
}

} // namespace ZShell