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

#include "segment.hpp"

#include "markdownparser.hpp"

#include <QCoreApplication>
#include <QDateTime>
#include <QPointer>
#include <QThreadPool>

namespace ZShell::llm {

namespace {
constexpr int kMarkdownRefreshMs = 150;
} // namespace

LlmSegment::LlmSegment(Type type, qint64 timestamp, QObject* parent)
	: QObject(parent), m_type(type), m_timestamp(timestamp) {
	m_markdownTimer.setInterval(kMarkdownRefreshMs);
	connect(&m_markdownTimer, &QTimer::timeout, this, [this]() {
		if (!m_markdownDirty) return;
		if (parseMarkdown()) m_markdownDirty = false;
	});
}

qint64 LlmSegment::elapsedMs() const {
	if (m_startedAt <= 0) return 0;
	const qint64 end = m_endedAt > 0 ? m_endedAt
									 : QDateTime::currentMSecsSinceEpoch();
	return end - m_startedAt;
}

void LlmSegment::begin() {
	if (m_running) return;
	m_running = true;
	m_startedAt = QDateTime::currentMSecsSinceEpoch();
	m_endedAt = 0;
	Q_EMIT runningChanged();
	Q_EMIT elapsedMsChanged();
}

void LlmSegment::close() {
	if (m_running) {
		m_running = false;
		m_endedAt = QDateTime::currentMSecsSinceEpoch();
		Q_EMIT runningChanged();
		Q_EMIT elapsedMsChanged();
	}
	if (m_type == Type::Content && m_markdownDirty) {
		m_markdownTimer.stop();
		if (parseMarkdown()) m_markdownDirty = false;
	}
}

void LlmSegment::appendText(const QString& piece) {
	if (piece.isEmpty()) return;
	m_text += piece;
	Q_EMIT textChanged();
	scheduleMarkdown();
}

void LlmSegment::setText(const QString& value) {
	if (m_text == value) return;
	m_text = value;
	Q_EMIT textChanged();
	scheduleMarkdown();
}

void LlmSegment::setName(const QString& value) {
	if (m_name == value) return;
	m_name = value;
	Q_EMIT nameChanged();
}

void LlmSegment::setToolCallId(const QString& value) {
	if (m_toolCallId == value) return;
	m_toolCallId = value;
	Q_EMIT toolCallIdChanged();
}

void LlmSegment::appendArguments(const QString& piece) {
	if (piece.isEmpty()) return;
	m_arguments += piece;
	Q_EMIT argumentsChanged();
}

void LlmSegment::setResult(const QString& value) {
	if (m_result == value) return;
	m_result = value;
	Q_EMIT resultChanged();
}

void LlmSegment::setStatus(Status value) {
	if (m_status == value) return;
	m_status = value;
	Q_EMIT statusChanged();
}

void LlmSegment::finishTool(const QString& resultText, bool success) {
	setResult(resultText);
	setStatus(success ? Status::Success : Status::Error);
	close();
}

void LlmSegment::restore(qint64 elapsedMs) {
	m_startedAt = m_timestamp;
	m_endedAt = m_timestamp + qMax<qint64>(0, elapsedMs);
}

void LlmSegment::scheduleMarkdown() {
	if (m_type != Type::Content) return;
	m_markdownDirty = true;
	if (!m_markdownTimer.isActive()) m_markdownTimer.start();
}

bool LlmSegment::parseMarkdown() {
	if (m_parseInFlight) return false;
	m_parseInFlight = true;
	const QString text = m_text;
	QThreadPool::globalInstance()->start([this, text]() {
		const QVariantList blocks = MarkdownParser::parse(text);
		QPointer<LlmSegment> guard(this);
		QMetaObject::invokeMethod(
			QCoreApplication::instance(),
			[guard, blocks]() {
				LlmSegment* seg = guard;
				if (!seg) return;
				seg->m_parseInFlight = false;
				seg->m_markdown = blocks;
				Q_EMIT seg->markdownChanged();
				if (seg->m_markdownDirty) seg->parseMarkdown();
			},
			Qt::QueuedConnection);
	});
	return true;
}

} // namespace ZShell::llm