#include "message.hpp"

#include "messagemodel.hpp"
#include "session.hpp"

namespace ZShell::llm {

namespace {

ChatSession* sessionOf(const ChatMessage* message) {
	if (auto* model = qobject_cast<ChatMessageModel*>(message->parent()))
		return model->session();
	return nullptr;
}

} // namespace

ChatMessage::ChatMessage(Role role, qint64 timestamp, QObject* parent)
: QObject(parent), m_role(role), m_timestamp(timestamp) {}

ChatGeneration* ChatMessage::addGeneration(qint64 timestamp) {
	auto* generation = new ChatGeneration(timestamp, this);
	m_generations.append(generation);
	if (m_active < 0)
		m_active = static_cast<int>(m_generations.size() - 1);
	Q_EMIT generationsChanged();
	return generation;
}

ChatGeneration* ChatMessage::appendGeneration(qint64 timestamp) {
	auto* generation = addGeneration(timestamp);
	setActiveInternal(static_cast<int>(m_generations.size() - 1));
	return generation;
}

void ChatMessage::removeGeneration(ChatGeneration* generation) {
	const int index = static_cast<int>(m_generations.indexOf(generation));
	if (index < 0)
		return;
	const bool wasActive = index == m_active;
	m_generations.removeAt(index);
	delete generation;
	if (m_generations.isEmpty()) {
		m_active = -1;
	} else if (m_active >= m_generations.size()) {
		m_active = static_cast<int>(m_generations.size() - 1);
	}
	Q_EMIT generationsChanged();
	if (wasActive)
		Q_EMIT activeGenerationChanged();
}

void ChatMessage::setActiveInternal(int index) {
	if (index < 0 || index >= m_generations.size() || index == m_active)
		return;
	m_active = index;
	Q_EMIT activeGenerationChanged();
}

void ChatMessage::setActiveGeneration(int index) {
	setActiveInternal(index);
}

void ChatMessage::edit(const QString& newContent) {
	if (auto* generation = activeGeneration())
		generation->setContent(newContent);
	if (auto* session = sessionOf(this))
		session->persist();
}

void ChatMessage::retry() {
	if (auto* session = sessionOf(this))
		session->retry(this);
}

void ChatMessage::generate() {
	if (auto* session = sessionOf(this))
		session->continueFrom(this);
}

} // namespace ZShell::llm