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

#include "messagemodel.hpp"

#include "session.hpp"

#include <algorithm>

namespace ZShell::llm {

ChatMessageModel::ChatMessageModel(ChatSession* session, QObject* parent)
	: QAbstractListModel(parent), m_session(session) {}

ChatMessageModel::~ChatMessageModel() = default;

int ChatMessageModel::rowCount(const QModelIndex& parent) const {
	if (parent.isValid()) return 0;
	return static_cast<int>(m_messages.size());
}

QVariant ChatMessageModel::data(const QModelIndex& index, int role) const {
	if (role != Qt::UserRole || !index.isValid() || index.row() < 0 ||
		index.row() >= m_messages.size())
		return {};
	return QVariant::fromValue(m_messages.at(index.row()));
}

QHash<int, QByteArray> ChatMessageModel::roleNames() const {
	return {{Qt::UserRole, "modelData"}};
}

ChatMessage* ChatMessageModel::at(int row) const {
	if (row < 0 || row >= m_messages.size()) return nullptr;
	return m_messages.at(row);
}

int ChatMessageModel::rowOf(const ChatMessage* message) const {
	for (int i = 0; i < m_messages.size(); ++i)
		if (m_messages.at(i) == message) return i;
	return -1;
}

ChatMessage* ChatMessageModel::createMessage(
	ChatMessage::Role role, qint64 timestamp) {
	return new ChatMessage(role, timestamp, this);
}

ChatMessage* ChatMessageModel::appendNewest(
	ChatMessage::Role role, const QString& content, qint64 timestamp) {
	auto* message = new ChatMessage(role, timestamp, this);
	auto* generation = message->addGeneration(timestamp);
	generation->setContent(content);

	beginInsertRows(QModelIndex(), 0, 0);
	m_messages.prepend(message);
	endInsertRows();

	emit lastMessageChanged();

	return message;
}

void ChatMessageModel::removeMessage(ChatMessage* message) {
	const int row = rowOf(message);
	if (row < 0) return;

	const bool wasLastMessage = row == 0;

	beginRemoveRows(QModelIndex(), row, row);
	m_messages.removeAt(row);
	endRemoveRows();

	delete message;

	if (wasLastMessage) emit lastMessageChanged();
}

void ChatMessageModel::removeRange(int firstRow, int lastRow) {
	if (firstRow < 0 || firstRow > lastRow || lastRow >= m_messages.size())
		return;

	const bool changesLastMessage = firstRow == 0;

	beginRemoveRows(QModelIndex(), firstRow, lastRow);
	for (int row = lastRow; row >= firstRow; --row)
		delete m_messages.takeAt(row);
	endRemoveRows();

	if (changesLastMessage) emit lastMessageChanged();
}

void ChatMessageModel::clear() {
	if (m_messages.isEmpty()) return;

	beginRemoveRows(QModelIndex(), 0, static_cast<int>(m_messages.size() - 1));
	qDeleteAll(m_messages);
	m_messages.clear();
	endRemoveRows();

	emit lastMessageChanged();
}

void ChatMessageModel::loadMessages(QList<ChatMessage*> messages) {
	beginResetModel();
	for (ChatMessage* message : m_messages)
		if (std::find(messages.begin(), messages.end(), message) ==
			messages.end())
			delete message;
	m_messages = std::move(messages);
	endResetModel();

	emit lastMessageChanged();
}

} // namespace ZShell::llm