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

#include "chatstore.hpp"

#include "chat.hpp"
#include "message.hpp"

#include <QDateTime>
#include <QDebug>
#include <QDir>
#include <QFileInfo>
#include <QStandardPaths>
#include <QSqlDatabase>
#include <QSqlError>
#include <QSqlQuery>
#include <QUuid>

#include <algorithm>

namespace ZShell {

ChatStore::ChatStore(QObject* parent)
: QObject(parent), m_connectionName(QUuid::createUuid().toString()) {
	openDb();
	load();
}

ChatStore::~ChatStore() {
	if (m_connectionName.isEmpty())
		return;
	QSqlDatabase db = QSqlDatabase::database(m_connectionName, false);
	db.close();
	QSqlDatabase::removeDatabase(m_connectionName);
}

QSqlDatabase ChatStore::db() const {
	return QSqlDatabase::database(m_connectionName);
}

void ChatStore::openDb() {
	const QString path =
		QStandardPaths::writableLocation(QStandardPaths::GenericStateLocation) +
		QStringLiteral("/zshell/chats.sqlite");
	QDir().mkpath(QFileInfo(path).absolutePath());

	QSqlDatabase db =
		QSqlDatabase::addDatabase(QStringLiteral("QSQLITE"), m_connectionName);
	db.setDatabaseName(path);
	if (!db.open()) {
		qWarning() << "ChatStore: failed to open database" << path << ":"
		           << db.lastError().text();
		return;
	}
	{
		QSqlQuery pragma(db);
		pragma.exec(QStringLiteral("PRAGMA foreign_keys = ON"));
	}
	{
		QSqlQuery query(db);
		query.exec(
			QStringLiteral(
				"CREATE TABLE IF NOT EXISTS sessions (\n"
				"  id TEXT PRIMARY KEY,\n"
				"  title TEXT NOT NULL DEFAULT '',\n"
				"  icon TEXT NOT NULL DEFAULT '',\n"
				"  created_at INTEGER NOT NULL,\n"
				"  updated_at INTEGER NOT NULL,\n"
				"  pinned INTEGER NOT NULL DEFAULT 0\n"
				")"));
	}
	{
		QSqlQuery info(db);
		bool hasIcon = false;
		if (info.exec(QStringLiteral("PRAGMA table_info(sessions)")))
			while (info.next())
				if (info.value(1).toString() == QLatin1String("icon")) {
					hasIcon = true;
					break;
				}
		if (!hasIcon) {
			QSqlQuery alter(db);
			if (!alter.exec(QStringLiteral(
					"ALTER TABLE sessions ADD COLUMN icon TEXT NOT NULL "
					"DEFAULT ''")))
				qWarning() << "ChatStore: failed to add icon column:"
				           << alter.lastError().text();
		}
	}
	{
		QSqlQuery query(db);
		query.exec(
			QStringLiteral(
		"CREATE TABLE IF NOT EXISTS messages (\n"
		"  session_id TEXT NOT NULL REFERENCES sessions (id) ON DELETE "
		"CASCADE,\n"
		"  role TEXT NOT NULL,\n"
		"  content TEXT,\n"
		"  reasoning TEXT,\n"
				"  timestamp INTEGER NOT NULL,\n"
				"  reasoning_elapsed_ms INTEGER NOT NULL DEFAULT 0,\n"
				"  content_elapsed_ms INTEGER NOT NULL DEFAULT 0\n"
				")"));
	}
}

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();
	const QString id = QString::number(now);
	{
		QSqlQuery query(db());
		query.prepare(
			"INSERT INTO sessions (id, title, created_at, updated_at) "
			"VALUES (:id, '', :created_at, :updated_at)");
		query.bindValue(":id", id);
		query.bindValue(":created_at", now);
		query.bindValue(":updated_at", now);
		if (!query.exec())
			qWarning() << "ChatStore: failed to insert session" << id << ":"
			           << query.lastError().text();
	}
	auto* session = new ChatSession(id, this);
	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 (auto* chat = qobject_cast<Chat*>(parent()))
		if (chat->m_active == session) {
			chat->stop();
			chat->endStream();
		}
	const QList<ChatSession*> before = m_sessions;
	{
		QSqlQuery query(db());
		query.prepare("DELETE FROM sessions WHERE id = :id");
		query.bindValue(":id", session->id());
		query.exec();
	}
	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*> sessions = m_sessions;
	for (ChatSession* session : sessions)
		removeSession(session);
}

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());
	session->ensureLoaded();
	if (!saveSession(session))
		return;
	sortAndNotify();
}

void ChatStore::saveMeta(ChatSession* session) {
	if (!session || !m_sessions.contains(session))
		return;
	QSqlQuery query(db());
	query.prepare("UPDATE sessions SET title = :title, icon = :icon "
	              "WHERE id = :id");
	query.bindValue(":title", session->title());
	query.bindValue(":icon", session->icon());
	query.bindValue(":id", session->id());
	if (!query.exec())
		qWarning() << "ChatStore: failed to save meta for" << session->id()
		           << ":" << query.lastError().text();
}

bool ChatStore::saveSession(ChatSession* session) {
	const QString id = session->id();
	QSqlDatabase handle = db();
	if (!handle.transaction()) {
		qWarning() << "ChatStore: failed to begin transaction:"
		           << handle.lastError().text();
		return false;
	}
	bool ok = true;
	{
		QSqlQuery query(handle);
		query.prepare(
			"UPDATE sessions SET title = :title, updated_at = :updated_at "
			"WHERE id = :id");
		query.bindValue(":title", session->title());
		query.bindValue(":updated_at", session->updatedAtMs());
		query.bindValue(":id", session->id());
		ok = query.exec();
	}
	if (ok) {
		QSqlQuery query(handle);
		query.prepare("DELETE FROM messages WHERE session_id = :id");
		query.bindValue(":id", session->id());
		ok = query.exec();
	}
	if (ok) {
		QSqlQuery insert(handle);
		ok = insert.prepare(
			"INSERT INTO messages (session_id, role, content, reasoning, "
			"timestamp, reasoning_elapsed_ms, content_elapsed_ms) "
			"VALUES (:id, :role, :content, :reasoning, :timestamp, "
			":reasoning_elapsed_ms, :content_elapsed_ms)");
		for (const auto* message : session->messages()) {
			insert.bindValue(":id", session->id());
			insert.bindValue(
				":role",
				message->role() == ChatMessage::Role::User
					? QStringLiteral("user")
					: QStringLiteral("assistant"));
			insert.bindValue(":content", message->content());
			insert.bindValue(":reasoning", message->reasoning());
			insert.bindValue(":timestamp", message->timestamp());
			insert.bindValue(":reasoning_elapsed_ms", message->reasoningElapsedMs());
			insert.bindValue(":content_elapsed_ms", message->contentElapsedMs());
			if (!insert.exec()) {
				ok = false;
				qWarning() << "ChatStore: saveSession" << id << "insert failed:"
				           << insert.lastError().text();
				break;
			}
		}
	}
	if (!ok || !handle.commit()) {
		qWarning() << "ChatStore: saveSession" << id << "commit failed, rolling back";
		handle.rollback();
		ok = false;
		qWarning() << "ChatStore: failed to save session" << session->id() << ":"
		           << handle.lastError().text();
	}
	return ok;
}

void ChatStore::loadMessagesInto(ChatSession* session) {
	QSqlQuery query(db());
	query.prepare(
		"SELECT role, content, reasoning, timestamp, reasoning_elapsed_ms, "
		"content_elapsed_ms FROM messages WHERE session_id = :id ORDER BY rowid");
	query.bindValue(":id", session->id());
	if (!query.exec()) {
		qWarning() << "ChatStore: failed to load messages for" << session->id()
		           << ":" << query.lastError().text();
		return;
	}
	QList<ChatMessage*> messages;
	while (query.next()) {
		auto* message = new ChatMessage(
			query.value(0).toString() == QLatin1String("user")
				? ChatMessage::Role::User
				: ChatMessage::Role::Assistant,
			query.value(1).toString(),
			query.value(3).toLongLong(),
			session);
		message->setReasoning(query.value(2).toString());
		message->setElapsedMs(
			query.value(4).toLongLong(), query.value(5).toLongLong());
		messages.append(message);
	}
	session->adoptMessages(messages);
}

void ChatStore::load() {
	QSqlQuery query(db());
	query.exec(
		"SELECT s.id, s.title, s.icon, s.created_at, s.updated_at, s.pinned, "
		"(SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) AS cnt "
		"FROM sessions s ORDER BY s.pinned DESC, s.updated_at DESC");
	while (query.next()) {
		auto* session = new ChatSession(query.value(0).toString(), this);
		session->setMeta(
			query.value(1).toString(),
			query.value(3).toLongLong(),
			query.value(4).toLongLong(),
			query.value(6).toInt());
		session->setIcon(query.value(2).toString());
		session->setPinned(query.value(5).toInt() != 0);
		m_sessions.append(session);
	}
	sortAndNotify();
}

void ChatStore::sortAndNotify() {
	const QList<ChatSession*> before = m_sessions;
	std::stable_sort(
		m_sessions.begin(),
		m_sessions.end(),
		[](const ChatSession* a, const ChatSession* b) {
			if (a->pinned() != b->pinned())
				return a->pinned() > b->pinned();
			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();
}

} // namespace ZShell