#include "message.hpp"

#include "session.hpp"

namespace ZShell::llm {

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

QString ChatMessage::content() const {
	if (m_role == Role::User) return m_content;
	const ChatGeneration* generation = activeGeneration();
	return generation ? generation->content() : QString();
}

void ChatMessage::setContent(const QString& value) {
	if (m_role != Role::User || m_content == value) return;
	m_content = value;
	Q_EMIT contentChanged();
}

ChatGeneration* ChatMessage::addGeneration(qint64 timestamp) {
	if (m_role != Role::Assistant) return nullptr;

	auto* generation = new ChatGeneration(this, timestamp, this);
	connect(
		generation, &ChatGeneration::contentChanged, this, [this, generation]() {
			if (generation == activeGeneration()) Q_EMIT contentChanged();
		});
	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) {
	ChatGeneration* 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();
		Q_EMIT contentChanged();
	}
}

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

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

void ChatMessage::setAttachments(QList<Attachment*> attachments) {
	QList<Attachment*> stale;
	for (Attachment* attachment : m_attachments)
		if (!attachments.contains(attachment)) stale.append(attachment);
	m_attachments.clear();
	for (Attachment* attachment : attachments) {
		attachment->setParent(this);
		m_attachments.append(attachment);
	}
	Q_EMIT attachmentsChanged();
	for (Attachment* attachment : stale)
		attachment->deleteLater();
}

void ChatMessage::edit(const QString& newContent) {
	if (m_role == Role::User) {
		setContent(newContent);
	} else if (ChatGeneration* generation = activeGeneration()) {
		generation->setContent(newContent);
	} else {
		return;
	}

	if (m_session) m_session->persist();
}

void ChatMessage::retry() {
	if (m_role != Role::Assistant || !m_session) return;
	m_session->retry(this);
}

void ChatMessage::generate() {
	if (m_role != Role::Assistant || !m_session) return;
	m_session->continueFrom(this);
}

} // namespace ZShell::llm