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) {}

LlmSegment::~LlmSegment() = default;

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();
	}
	closed();
}

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

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

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

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

void LlmTextSegment::closed() {
	if (type() != Type::Content) return;
	if (m_markdownDirty) {
		m_markdownTimer.stop();
		if (parseMarkdown()) m_markdownDirty = false;
	}
}

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

bool LlmTextSegment::parseMarkdown() {
	if (m_parseInFlight) return false;
	m_parseInFlight = true;
	const QString text = m_text;
	QThreadPool::globalInstance()->start([this, text]() {
		const QVector<MarkdownBlockData> parsed = MarkdownParser::parse(text);
		QPointer<LlmTextSegment> guard(this);
		QMetaObject::invokeMethod(
			QCoreApplication::instance(),
			[guard, parsed]() {
				LlmTextSegment* segment = guard;
				if (!segment) return;
				segment->m_parseInFlight = false;

				QList<MarkdownBlock*> blocks;
				blocks.reserve(static_cast<int>(parsed.size()));
				for (int i = 0; i < static_cast<int>(parsed.size()); ++i) {
					const MarkdownBlockData& data = parsed.at(i);
					const QString id = QStringLiteral("%1:%2").arg(i).arg(
						static_cast<int>(data.type));
					blocks.append(new MarkdownBlock(
						data.type,
						id,
						data.text,
						data.latex,
						data.code,
						data.language,
						data.level,
						segment));
				}
				segment->m_blocks = blocks;
				Q_EMIT segment->markdownChanged();
				if (segment->m_markdownDirty) segment->parseMarkdown();
			},
			Qt::QueuedConnection);
	});
	return true;
}

LlmToolCall::LlmToolCall(qint64 timestamp, QObject* parent)
	: LlmSegment(Type::ToolCall, timestamp, parent) {}

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

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

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

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

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

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

} // namespace ZShell::llm