#include "llmclient.hpp"

#include "attachment.hpp"
#include "generation.hpp"
#include "message.hpp"
#include "messagemodel.hpp"
#include "segment.hpp"
#include "session.hpp"

#include <QDebug>
#include <QFile>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QNetworkReply>
#include <QNetworkRequest>
#include <QRegularExpression>
#include <QSet>
#include <QUrl>

#include <memory>

namespace ZShell::llm {

namespace {

constexpr int kMinContextTokens = 512;

qsizetype jsonValueChars(const QJsonValue& value) {
	switch (value.type()) {
	case QJsonValue::String:
		return value.toString().size();
	case QJsonValue::Array: {
		qsizetype total = 0;
		for (const QJsonValue& item : value.toArray())
			total += jsonValueChars(item);
		return total;
	}
	case QJsonValue::Object: {
		qsizetype total = 0;
		const QJsonObject obj = value.toObject();
		for (auto it = obj.constBegin(); it != obj.constEnd(); ++it)
			total += jsonValueChars(it.value());
		return total;
	}
	default:
		return 2;
	}
}

int estimateMessageTokens(const QJsonObject& message) {
	return int(jsonValueChars(message) / LlmTool::CharsPerToken) + 4;
}

struct ContextUnit {
	QList<QJsonObject> messages;
	int tokens = 0;
};

void shrinkUnitToBudget(QList<QJsonObject>& unit, int budgetTokens) {
	static const QStringList kTextKeys = {
		QStringLiteral("content"), QStringLiteral("reasoning_content")};
	for (int pass = 0; pass < 8; ++pass) {
		int tokens = 0;
		for (const QJsonObject& message : unit)
			tokens += estimateMessageTokens(message);
		if (tokens <= budgetTokens) return;
		qsizetype bestMessage = -1;
		QString bestKey;
		qsizetype bestLength = 0;
		for (qsizetype i = 0; i < unit.size(); ++i) {
			const QJsonObject message = unit.at(i);
			for (const QString& key : kTextKeys) {
				const QJsonValue value = message.value(key);
				if (value.isString() && value.toString().size() > bestLength) {
					bestMessage = i;
					bestKey = key;
					bestLength = value.toString().size();
				}
			}
		}
		if (bestMessage < 0) return;
		QString text = unit.at(bestMessage).value(bestKey).toString();
		text.truncate(qMax<qsizetype>(32, text.size() * 3 / 4));
		const QString marker = QStringLiteral("\n[…]");
		if (!text.endsWith(marker)) text += marker;
		unit[bestMessage][bestKey] = text;
	}
}

int estimateArrayTokens(const QJsonArray& messages) {
	int total = 0;
	for (const QJsonValue& value : messages)
		total += estimateMessageTokens(value.toObject());
	return total;
}

constexpr int kToolResultStubChars = 512;

constexpr double kPromptEstimateDiscount = 10.0 / 12.0;

constexpr int kMinToolResultTokens = 128;

int completionReserve(int contextSize) {
	return qBound(2048, contextSize / 8, 8192);
}

int promptBudgetTokens(int contextSize) {
	return int(
		(contextSize - completionReserve(contextSize)) *
		kPromptEstimateDiscount);
}

constexpr int kRefreshIntervalMs = 60'000;

constexpr int kProbeTimeoutMs = 2000;

int contextSizeFromErrorMessage(const QString& message) {
	static const QRegularExpression re(
		QLatin1String("available context size \\((\\d+)"));
	const auto match = re.match(message);
	if (match.hasMatch()) return match.captured(1).toInt();
	return 0;
}

} // namespace

QString LlmClient::completionsPath(
	const QString& endpoint, const QString& subpath) {
	QString base = endpoint.trimmed();
	while (base.endsWith('/'))
		base.chop(1);
	if (!base.endsWith("/v1")) base += "/v1";
	return base + subpath;
}

QString LlmClient::jsonErrorMessage(const QByteArray& body) {
	const QJsonDocument doc = QJsonDocument::fromJson(body);
	if (!doc.isObject()) return QString();
	const QJsonValue errorValue = doc.object().value("error");
	if (errorValue.isObject())
		return errorValue.toObject().value("message").toString();
	if (errorValue.isString()) return errorValue.toString();
	return QString();
}

QString LlmClient::serverErrorMessage(
	const QByteArray& body, const QString& fallback) {
	const QString message = jsonErrorMessage(body);
	if (!message.isEmpty()) return message;
	return fallback.isEmpty() ? QStringLiteral("Request to LLM server failed")
							  : fallback;
}

void LlmClient::setServerError(const QString& value) {
	if (m_serverError == value) return;
	m_serverError = value;
	Q_EMIT serverErrorChanged();
}

void LlmClient::setBusy(bool value) {
	if (m_busy == value) return;
	m_busy = value;
	Q_EMIT busyChanged();
}

void LlmClient::setStreamingChatId(const QString& id) {
	if (m_streamingChatId == id) return;
	m_streamingChatId = id;
	Q_EMIT streamingChatIdChanged();
}

LlmClient::LlmClient(QObject* parent) : QObject(parent) {
	m_tools = new ToolRegistry(this);
	connect(
		m_tools,
		&ToolRegistry::enabledChanged,
		this,
		&LlmClient::toolsEnabledChanged);

	m_refreshTimer.setParent(this);
	m_refreshTimer.setInterval(kRefreshIntervalMs);
	connect(
		&m_refreshTimer, &QTimer::timeout, this, &LlmClient::refreshAllStates);
	m_refreshTimer.start();
}

LlmClient::~LlmClient() {
	if (m_reply) {
		QNetworkReply* reply = m_reply;
		m_reply = nullptr;
		reply->abort();
	}
	m_tools->cancelAll();
}

void LlmClient::setEndpoint(const QString& value) {
	if (m_endpoint == value) return;
	m_endpoint = value;
	Q_EMIT endpointChanged();
	m_refreshTimer.start();
	refreshAllStates();
}

void LlmClient::setModel(const QString& value) {
	if (m_model == value) return;
	m_model = value;
	Q_EMIT modelChanged();
	if (m_model.isEmpty()) refreshModels();
}

void LlmClient::setTemperature(double value) {
	m_temperature = value;
}

void LlmClient::setSandboxEnabled(bool value) {
	m_sandboxEnabled = value;
}

QString LlmClient::activeWorkspace() const {
	return m_active ? m_active->workspacePath() : QString();
}

void LlmClient::startGeneration(ChatSession* session, ChatGeneration* target) {
	if (m_busy || !session || !target) return;
	m_contextRetryUsed = false;
	m_active = session;
	m_streaming = target;
	target->setStreaming(true);
	setBusy(true);
	setStreamingChatId(session->id());

	const QUrl url =
		QUrl::fromUserInput(completionsPath(m_endpoint, "/chat/completions"));
	if (!url.isValid() || url.host().isEmpty()) {
		fail(QStringLiteral("Invalid LLM endpoint: %1").arg(m_endpoint));
		return;
	}

	const auto* model = session->messagesModel();
	const int targetRow =
		model->rowOf(qobject_cast<ChatMessage*>(target->parent()));
	if (targetRow < 0) {
		fail(QStringLiteral(
			"Internal error: generation target is not in the "
			"session"));
		return;
	}

	m_transcript = QJsonArray();
	m_callBuilders.clear();
	m_callResults.clear();
	m_finishReason.clear();
	m_round = 0;
	m_contentMark = 0;
	m_reasoningMark = 0;
	setApprovalPending(false);

	probeContextSize([this, session, target]() {
		if (m_active != session || m_streaming != target) return;
		refreshModels();
		sendRound();
	});
}

void LlmClient::sendRound() {
	if (!m_active || !m_streaming) return;
	probeContextSize();
	m_finishReason.clear();
	m_callBuilders.clear();
	m_callResults.clear();
	m_roundDone = false;
	m_streamDone = false;
	m_toolPhase = false;
	m_pendingCalls = 0;
	m_buffer.clear();

	const QUrl url =
		QUrl::fromUserInput(completionsPath(m_endpoint, "/chat/completions"));
	if (!url.isValid() || url.host().isEmpty()) {
		fail(QStringLiteral("Invalid LLM endpoint: %1").arg(m_endpoint));
		return;
	}

	ChatSession* active = m_active;
	ChatGeneration* streaming = m_streaming;
	const auto* model = active->messagesModel();
	const int targetRow =
		model->rowOf(qobject_cast<ChatMessage*>(streaming->parent()));
	if (targetRow < 0) {
		fail(QStringLiteral(
			"Internal error: generation target is not in the "
			"session"));
		return;
	}

	int transcriptTokens = 0;
	for (const QJsonValue& value : m_transcript)
		transcriptTokens += estimateMessageTokens(value.toObject());
	QJsonArray messages =
		buildContextMessages(m_active, targetRow, transcriptTokens);
	for (const QJsonValue& value : m_transcript)
		messages.append(value);

	if (m_contextSize > 0) {
		const int promptLimit = promptBudgetTokens(m_contextSize);
		while (estimateArrayTokens(messages) > promptLimit) {
			bool stubbed = false;
			for (qsizetype i = 0; i < messages.size(); ++i) {
				QJsonObject message = messages.at(i).toObject();
				if (message.value("role").toString() != QLatin1String("tool"))
					continue;
				const QString content = message.value("content").toString();
				if (content.size() <= kToolResultStubChars) continue;
				QString stub = content.left(kToolResultStubChars);
				stub += QStringLiteral(
							"\n[truncated to fit the context window; was "
							"%1 characters]")
							.arg(content.size());
				message["content"] = stub;
				messages[i] = message;
				stubbed = true;
				break;
			}
			if (!stubbed) break;
		}
	}

	appendSandboxPreamble(messages);

	QJsonObject body;
	QJsonObject streamOptions;
	streamOptions[QStringLiteral("include_usage")] = true;
	body[QStringLiteral("stream_options")] = streamOptions;
	body[QStringLiteral("messages")] = messages;
	body[QStringLiteral("stream")] = true;
	body[QStringLiteral("temperature")] = m_temperature;
	if (!m_model.isEmpty()) body[QStringLiteral("model")] = m_model;
	const QJsonArray toolSpecs = m_tools->specifications();
	if (!toolSpecs.isEmpty()) body[QStringLiteral("tools")] = toolSpecs;

	QNetworkRequest request(url);
	request.setHeader(QNetworkRequest::ContentTypeHeader, "application/json");
	request.setRawHeader("Accept", "text/event-stream");

	QNetworkReply* reply =
		m_manager.post(request, QJsonDocument(body).toJson());
	if (!reply) {
		fail(QStringLiteral("Failed to start request to LLM server"));
		return;
	}
	m_reply = reply;

	connect(reply, &QNetworkReply::readyRead, this, [this, reply]() {
		if (m_reply != reply) return;
		m_buffer.append(reply->readAll());
		drainBuffer();
	});
	connect(reply, &QNetworkReply::finished, this, [this, reply]() {
		if (m_reply != reply) {
			reply->deleteLater();
			return;
		}
		m_reply = nullptr;

		const QNetworkReply::NetworkError error = reply->error();
		const QString errorString = reply->errorString();
		const QByteArray responseBody = reply->readAll();
		m_buffer.append(responseBody);
		reply->deleteLater();

		drainBuffer();
		if (!m_streaming) return;

		if (error == QNetworkReply::NoError) {
			if (m_streamDone) {
				roundFinished();
				return;
			}
			const QByteArray tail = m_buffer;
			m_buffer.clear();
			const QString serverMessage = jsonErrorMessage(tail.simplified());
			fail(
				serverMessage.isEmpty()
					? QStringLiteral(
						  "LLM server closed the stream before "
						  "completing the response")
					: serverMessage);
			return;
		}
		if (error == QNetworkReply::OperationCanceledError) {
			finishTurn();
			return;
		}
		const QString message = serverErrorMessage(responseBody, errorString);
		const int reported = contextSizeFromErrorMessage(message);
		if (reported > 0 && reported != m_contextSize && !m_contextRetryUsed) {
			qWarning() << "LlmClient: prompt overflow, server reports"
					   << reported << "tokens; retrying round";
			setContextSize(reported);
			m_contextRetryUsed = true;
			sendRound();
			return;
		}
		fail(message);
	});
}

QJsonArray LlmClient::buildContextMessages(
	ChatSession* session, int stopBeforeRow, int extraReserveTokens) const {
	const auto* model = session->messagesModel();
	QList<ContextUnit> units;
	for (int row = model->rowCount() - 1; row > stopBeforeRow; --row) {
		const auto* message = model->at(row);
		const auto* generation = message->activeGeneration();
		if (!generation) continue;
		const bool hasAnyContent = !generation->content().isEmpty() ||
								   !generation->reasoning().isEmpty() ||
								   generation->toolCallCount() > 0;
		if (message->role() != ChatMessage::Role::User && !hasAnyContent)
			continue;

		ContextUnit unit;
		if (message->role() == ChatMessage::Role::User) {
			QJsonObject user;
			user[QStringLiteral("role")] = QStringLiteral("user");
			const QString text = generation->content();
			const QList<Attachment*> attachments = message->attachments();
			if (attachments.isEmpty()) {
				user[QStringLiteral("content")] = text;
			} else {
				QJsonArray parts;
				if (!text.isEmpty()) {
					QJsonObject textPart;
					textPart[QStringLiteral("type")] = QStringLiteral("text");
					textPart[QStringLiteral("text")] = text;
					parts.append(textPart);
				}
				for (const auto* attachment : attachments) {
					if (attachment->kind() != Attachment::Kind::Image) continue;
					QFile file(attachment->path());
					if (!file.open(QIODevice::ReadOnly)) continue;
					const QByteArray base64 = file.readAll().toBase64();
					file.close();
					if (base64.isEmpty()) continue;
					QJsonObject imageUrl;
					imageUrl[QStringLiteral("url")] =
						QStringLiteral("data:%1;base64,%2")
							.arg(attachment->mimeType())
							.arg(QString::fromLatin1(base64));
					QJsonObject imagePart;
					imagePart[QStringLiteral("type")] =
						QLatin1String("image_url");
					imagePart[QStringLiteral("image_url")] = imageUrl;
					parts.append(imagePart);
				}
				if (parts.isEmpty())
					user[QStringLiteral("content")] = text;
				else
					user[QStringLiteral("content")] = parts;
			}
			unit.messages.append(user);
		} else {
			QList<const LlmSegment*> toolSegments;
			for (const auto* segment : generation->segments())
				if (segment->type() == LlmSegment::Type::ToolCall)
					toolSegments.append(segment);

			QJsonObject assistant;
			assistant[QStringLiteral("role")] = QStringLiteral("assistant");
			if (toolSegments.isEmpty())
				assistant[QStringLiteral("content")] = generation->content();
			else if (!generation->content().isEmpty())
				assistant[QStringLiteral("content")] = generation->content();
			if (!generation->reasoning().isEmpty())
				assistant[QStringLiteral("reasoning_content")] =
					generation->reasoning();
			if (!toolSegments.isEmpty()) {
				QJsonArray calls;
				for (const auto* segment : toolSegments) {
					QJsonObject function;
					function[QStringLiteral("name")] = segment->name();
					function[QStringLiteral("arguments")] =
						segment->arguments();
					QJsonObject call;
					call[QStringLiteral("id")] = segment->toolCallId();
					call[QStringLiteral("type")] = QStringLiteral("function");
					call[QStringLiteral("function")] = function;
					calls.append(call);
				}
				assistant[QStringLiteral("tool_calls")] = calls;
			}
			unit.messages.append(assistant);
			for (const auto* segment : toolSegments) {
				QJsonObject toolMessage;
				toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
				toolMessage[QStringLiteral("tool_call_id")] =
					segment->toolCallId();
				toolMessage[QStringLiteral("content")] = segment->result();
				unit.messages.append(toolMessage);
			}
		}
		for (const QJsonObject& messageJson : unit.messages)
			unit.tokens += estimateMessageTokens(messageJson);
		units.append(unit);
	}

	if (m_contextSize > 0) {
		int budget = m_contextSize - completionReserve(m_contextSize) -
					 extraReserveTokens;
		budget = qMax(budget, kMinContextTokens);

		int total = 0;
		for (const ContextUnit& unit : units)
			total += unit.tokens;
		while (units.size() > 1 && total > budget) {
			total -= units.first().tokens;
			units.removeFirst();
		}
		while (
			units.size() > 1 &&
			units.first().messages.first()[QStringLiteral("role")].toString() ==
				QLatin1String("assistant")) {
			total -= units.first().tokens;
			units.removeFirst();
		}
		if (units.size() == 1 && units.first().tokens > budget)
			shrinkUnitToBudget(units.first().messages, budget);
	}

	QJsonArray messages;
	for (const ContextUnit& unit : units)
		for (const QJsonObject& messageJson : unit.messages)
			messages.append(messageJson);
	return messages;
}

void LlmClient::appendSandboxPreamble(QJsonArray& messages) {
	if (!m_sandboxEnabled) return;
	ChatSession* active = m_active;
	if (!active) return;
	const QString workspace = active->workspacePath();
	if (workspace.isEmpty()) return;
	const auto* model = active->messagesModel();
	QStringList files;
	QSet<QString> seenPaths;
	for (int row = 0; row < model->rowCount(); ++row) {
		const auto* message = model->at(row);
		if (message->role() != ChatMessage::Role::User) continue;
		for (const auto* attachment : message->attachments()) {
			if (attachment->kind() != Attachment::Kind::File) continue;
			const QString relativePath =
				attachment->path().mid(workspace.size() + 1);
			if (seenPaths.contains(relativePath)) continue;
			seenPaths.insert(relativePath);
			files.append(QStringLiteral("- %1 (available as %2)")
							 .arg(attachment->name(), relativePath));
		}
	}
	if (files.isEmpty()) return;
	QJsonObject preamble;
	preamble[QStringLiteral("role")] = QStringLiteral("system");
	preamble[QStringLiteral("content")] =
		QStringLiteral(
			"The user attached the following files to this "
			"conversation. They are available in your sandbox "
			"directory at %1:\n%2\nruncommand executes shell "
			"commands inside a bubblewrap sandbox: the rest of the "
			"filesystem is read-only, there is no network access, "
			"and only the sandbox directory is writable. Use it to "
			"inspect, transform, or extract attached files (for "
			"example to extract archives or search across many "
			"files). Files you create stay inside the sandbox "
			"directory. Use the webfetch tool for anything that "
			"needs the network.")
			.arg(workspace, files.join(QStringLiteral("\n")));
	messages.prepend(preamble);
}

void LlmClient::applyToolCallDelta(const QJsonObject& call) {
	ChatGeneration* streaming = m_streaming;
	if (!streaming) return;
	const int index = call[QStringLiteral("index")].toInt(-1);
	if (index < 0) return;
	while (m_callBuilders.size() <= index)
		m_callBuilders.append(ToolCallBuilder{});
	ToolCallBuilder& builder = m_callBuilders[index];
	builder.seen = true;
	const QString id = call[QStringLiteral("id")].toString();
	if (!id.isEmpty()) builder.id = id;
	const QJsonObject function = call[QStringLiteral("function")].toObject();
	const QString name = function[QStringLiteral("name")].toString();
	if (!name.isEmpty()) builder.name = name;
	const QString arguments = function[QStringLiteral("arguments")].toString();
	if (!arguments.isEmpty()) builder.arguments += arguments;

	if (!builder.segment) {
		streaming->closeOpenSegments();
		builder.segment = streaming->beginToolCall(builder.name, builder.id);
	}
	builder.segment->setName(builder.name);
	builder.segment->setToolCallId(builder.id);
	if (!arguments.isEmpty()) builder.segment->appendArguments(arguments);
}

void LlmClient::roundFinished() {
	if (m_roundDone || !m_streaming) return;
	ChatGeneration* streaming = m_streaming;
	m_roundDone = true;

	bool hasCalls = false;
	for (const auto& builder : m_callBuilders)
		if (builder.seen) hasCalls = true;
	if (!hasCalls) {
		finishTurn();
		return;
	}
	if (m_finishReason != QLatin1String("tool_calls")) {
		qWarning() << "LlmClient: stream ended" << m_finishReason
				   << "with incomplete tool calls, ending turn";
		for (const auto& builder : m_callBuilders) {
			if (!builder.seen || !builder.segment) continue;
			builder.segment->finishTool(
				QStringLiteral(
					"Stream truncated before the tool call "
					"completed"),
				false);
		}
		finishTurn();
		return;
	}
	if (m_round >= kMaxToolRounds) {
		qWarning() << "LlmClient: tool round limit reached, ending turn";
		finishTurn();
		return;
	}

	streaming->closeOpenSegments();

	const QString content = streaming->content();
	const QString reasoning = streaming->reasoning();
	QJsonObject assistant;
	assistant[QStringLiteral("role")] = QStringLiteral("assistant");
	if (content.size() > m_contentMark) {
		assistant[QStringLiteral("content")] = content.mid(m_contentMark);
		m_contentMark = content.size();
	}
	if (reasoning.size() > m_reasoningMark) {
		assistant[QStringLiteral("reasoning_content")] =
			reasoning.mid(m_reasoningMark);
		m_reasoningMark = reasoning.size();
	}
	QJsonArray calls;
	for (const auto& builder : m_callBuilders) {
		if (!builder.seen) continue;
		QJsonObject function;
		function[QStringLiteral("name")] = builder.name;
		function[QStringLiteral("arguments")] = builder.arguments;
		QJsonObject call;
		call[QStringLiteral("id")] = builder.id;
		call[QStringLiteral("type")] = QStringLiteral("function");
		call[QStringLiteral("function")] = function;
		calls.append(call);
	}
	assistant[QStringLiteral("tool_calls")] = calls;
	m_transcript.append(assistant);
	m_round++;

	setApprovalPending(true);
}

void LlmClient::executeAllCalls() {
	m_toolPhase = true;
	m_pendingCalls = 0;
	m_callResults = QList<ToolCallResult>(m_callBuilders.size());

	int batchSize = 0;
	for (const auto& call : m_callBuilders)
		if (call.seen) ++batchSize;
	int perCallChars = 0;
	if (batchSize > 0 && m_contextSize > 0) {
		const int available = promptBudgetTokens(m_contextSize) -
							  estimateArrayTokens(m_transcript);
		const int perCallTokens =
			qMax(available / batchSize, kMinToolResultTokens);
		perCallChars = perCallTokens * LlmTool::CharsPerToken;
	}

	for (int i = 0; i < m_callBuilders.size(); ++i) {
		const auto& call = m_callBuilders.at(i);
		if (!call.seen) continue;

		LlmTool* tool = m_tools->tool(call.name);
		QJsonObject args;
		QString errorText;
		if (!tool) {
			errorText =
				QStringLiteral("Error: unknown tool '%1'").arg(call.name);
		} else if (!call.arguments.isEmpty()) {
			const QJsonDocument doc =
				QJsonDocument::fromJson(call.arguments.toUtf8());
			if (!doc.isObject()) {
				errorText = QStringLiteral(
								"Error: tool arguments are not valid "
								"JSON: %1")
								.arg(call.arguments);
			} else {
				args = doc.object();
			}
		}
		if (!errorText.isEmpty()) {
			m_callResults[i] = {errorText, false};
			if (LlmSegment* segment = call.segment)
				segment->finishTool(errorText, false);
			continue;
		}

		++m_pendingCalls;
		tool->execute(
			call.id,
			args,
			perCallChars,
			[this, i, call](const QJsonObject& result) {
				if (!m_streaming) return;
				const bool success = result.contains(QStringLiteral("output"));
				const QString content =
					success ? result[QStringLiteral("output")].toString()
							: QStringLiteral("Error: ") +
								  result[QStringLiteral("error")].toString();
				m_callResults[i] = {content, success};
				if (LlmSegment* segment = call.segment)
					segment->finishTool(content, success);
				if (--m_pendingCalls == 0) flushCallResults();
			});
	}
	if (m_pendingCalls == 0) flushCallResults();
}

void LlmClient::flushCallResults() {
	if (!m_toolPhase) return;
	m_toolPhase = false;
	if (!m_streaming) return;
	for (int i = 0; i < m_callBuilders.size(); ++i) {
		const auto& call = m_callBuilders.at(i);
		if (!call.seen) continue;
		QJsonObject toolMessage;
		toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
		toolMessage[QStringLiteral("tool_call_id")] = call.id;
		toolMessage[QStringLiteral("content")] = m_callResults.at(i).content;
		m_transcript.append(toolMessage);
	}
	sendRound();
}

void LlmClient::setApprovalPending(bool value) {
	if (m_approvalPending == value) return;
	m_approvalPending = value;
	if (ChatGeneration* streaming = m_streaming)
		streaming->setApprovalPending(value);
}

void LlmClient::approveTools() {
	if (!m_approvalPending) return;
	ChatGeneration* streaming = m_streaming;
	if (!streaming) return;
	setApprovalPending(false);
	for (auto* segment : streaming->segments()) {
		if (segment->type() != LlmSegment::Type::ToolCall ||
			segment->status() != LlmSegment::Status::Pending)
			continue;
		segment->setStatus(LlmSegment::Status::Running);
		segment->begin();
	}
	executeAllCalls();
}

void LlmClient::denyTools() {
	if (!m_approvalPending || !m_streaming) return;
	setApprovalPending(false);
	finishPendingCalls(QStringLiteral("Denied by user"));
	finishTurn();
}

void LlmClient::finishPendingCalls(const QString& resultText) {
	ChatGeneration* streaming = m_streaming;
	if (!streaming) return;
	for (auto* segment : streaming->segments())
		if (segment->type() == LlmSegment::Type::ToolCall &&
			segment->status() == LlmSegment::Status::Pending)
			segment->finishTool(resultText, false);
}

void LlmClient::stop() {
	if (!m_busy) return;
	if (m_approvalPending) {
		setApprovalPending(false);
		finishPendingCalls(QStringLiteral("Cancelled"));
		finishTurn();
		return;
	}
	if (m_toolPhase) {
		m_tools->cancelAll();
		if (ChatGeneration* streaming = m_streaming) {
			for (auto* segment : streaming->segments()) {
				if (segment->type() == LlmSegment::Type::ToolCall &&
					segment->running())
					segment->finishTool(QStringLiteral("Cancelled"), false);
			}
		}
		finishTurn();
		return;
	}
	if (m_reply) {
		QNetworkReply* reply = m_reply;
		m_reply = nullptr;
		reply->abort();
	}
	endStream();
}

void LlmClient::endStream() {
	if (!m_streaming) return;
	ChatGeneration* generation = m_streaming;
	ChatSession* session = m_active;
	m_streaming = nullptr;
	m_active = nullptr;
	if (m_reply) {
		QNetworkReply* reply = m_reply;
		m_reply = nullptr;
		reply->abort();
	}
	generation->setStreaming(false);
	m_approvalPending = false;
	generation->setApprovalPending(false);
	for (auto* segment : generation->segments())
		if (segment->type() == LlmSegment::Type::ToolCall &&
			segment->status() == LlmSegment::Status::Pending)
			segment->finishTool(QStringLiteral("Cancelled"), false);
	if (generation->content().isEmpty() && generation->reasoning().isEmpty() &&
		generation->toolCallCount() == 0 && !generation->error()) {
		if (auto* message = qobject_cast<ChatMessage*>(generation->parent())) {
			if (message->generationCount() <= 1) {
				if (session) session->removeMessage(message);
			} else {
				message->removeGeneration(generation);
			}
		}
	}
	setBusy(false);
	setStreamingChatId(QString());
}

void LlmClient::finishTurn() {
	if (!m_streaming) return;
	ChatSession* session = m_active;
	endStream();
	if (session && m_pendingClear == session) {
		session->clearMessages();
		m_pendingClear.clear();
	}
	if (session) session->persist();
	refreshAllStates();
}

void LlmClient::clearOnFinish(ChatSession* session) {
	m_pendingClear = session;
}

void LlmClient::sessionRemoved(ChatSession* session) {
	if (m_pendingClear == session) m_pendingClear.clear();
	if (m_active == session) {
		stop();
		endStream();
	}
}

void LlmClient::fail(const QString& message) {
	qWarning() << "LlmClient:" << message;
	if (ChatGeneration* streaming = m_streaming) streaming->setError(message);
	finishTurn();
	Q_EMIT errorOccurred(message);
}

void LlmClient::drainBuffer() {
	while (true) {
		const qsizetype newline = m_buffer.indexOf('\n');
		if (newline < 0) break;
		const QByteArray line = m_buffer.left(newline).trimmed();
		m_buffer.remove(0, newline + 1);
		handleLine(line);
	}
}

void LlmClient::handleLine(const QByteArray& line) {
	if (!m_streaming || line.isEmpty() || !line.startsWith("data:")) return;

	const QByteArray data = line.mid(5).trimmed();
	if (data == "[DONE]") {
		m_streamDone = true;
		roundFinished();
		return;
	}

	const QJsonDocument doc = QJsonDocument::fromJson(data);
	if (!doc.isObject()) return;
	const QJsonObject obj = doc.object();
	updateTokenUsage(obj);

	if (obj.contains("error")) {
		const QJsonValue errorValue = obj["error"];
		QString message;
		if (errorValue.isObject())
			message = errorValue.toObject()["message"].toString();
		else if (errorValue.isString())
			message = errorValue.toString();
		fail(
			message.isEmpty() ? QStringLiteral("LLM server returned an error")
							  : message);
		return;
	}

	for (const QJsonValue& choiceValue : obj["choices"].toArray()) {
		if (!m_streaming) continue;
		const QJsonObject choice = choiceValue.toObject();
		const QJsonObject delta = choice["delta"].toObject();

		const QString finishReason =
			choice[QStringLiteral("finish_reason")].toString();
		if (!finishReason.isEmpty()) m_finishReason = finishReason;

		m_streaming->appendContent(delta["content"].toString());
		QString reasoning = delta["reasoning_content"].toString();
		if (reasoning.isEmpty()) reasoning = delta["reasoning"].toString();
		m_streaming->appendReasoning(reasoning);

		for (const QJsonValue& callValue : delta["tool_calls"].toArray()) {
			if (!m_streaming) break;
			applyToolCallDelta(callValue.toObject());
		}
	}
}

void LlmClient::refreshAllStates() {
	if (m_endpoint.trimmed().isEmpty()) {
		setServerError(QStringLiteral("LLM endpoint is not configured"));
		return;
	}
	probeContextSize([this]() { refreshModels(); });
}

void LlmClient::refreshModels() {
	const QUrl url =
		QUrl::fromUserInput(completionsPath(m_endpoint, "/models"));
	if (!url.isValid() || url.host().isEmpty()) {
		setServerError(QStringLiteral("LLM endpoint is not configured"));
		return;
	}

	auto* reply = m_manager.get(QNetworkRequest(url));
	connect(reply, &QNetworkReply::finished, this, [this, reply]() {
		const QNetworkReply::NetworkError error = reply->error();
		const QString errorString = reply->errorString();
		const QByteArray data = reply->readAll();
		reply->deleteLater();

		if (error != QNetworkReply::NoError) {
			const QString serverMessage = jsonErrorMessage(data);
			setServerError(
				serverMessage.isEmpty()
					? QStringLiteral("Cannot reach LLM server: %1")
						  .arg(errorString)
					: serverMessage);
			return;
		}
		const QJsonArray dataArr =
			QJsonDocument::fromJson(data).object()["data"].toArray();

		QStringList models;
		QSet<QString> seen;
		for (const QJsonValue& value : dataArr) {
			const QString id = value.toObject()["id"].toString();
			if (!id.isEmpty() && !seen.contains(id)) {
				seen.insert(id);
				models.append(id);
			}
		}
		if (models.isEmpty()) {
			setServerError(QStringLiteral(
				"LLM server is running but has no models loaded"));
			return;
		}
		setServerError(QString());

		m_availableModels = models;
		Q_EMIT availableModelsChanged();

		if (m_model.isEmpty()) {
			m_model = models.first();
			Q_EMIT modelChanged();
		} else if (!models.contains(m_model)) {
			m_model = models.first();
			Q_EMIT modelChanged();
		}
	});
}

void LlmClient::setContextSize(int size) {
	if (size <= 0 || m_contextSize == size) return;
	m_contextSize = size;
	m_tools->setContextSize(size);
	Q_EMIT contextSizeChanged();
}

void LlmClient::probeContextSize(std::function<void()> done) {
	QString base = m_endpoint.trimmed();
	while (base.endsWith('/'))
		base.chop(1);
	const QUrl url = QUrl::fromUserInput(base + "/props");
	if (!url.isValid() || url.host().isEmpty()) {
		setServerError(
			QStringLiteral("Invalid LLM endpoint: %1").arg(m_endpoint));
		if (done) done();
		return;
	}

	auto* reply = m_manager.get(QNetworkRequest(url));
	auto settled = std::make_shared<bool>(false);
	const auto finish = [this, reply, settled, done]() {
		if (*settled) return;
		*settled = true;

		const QNetworkReply::NetworkError error = reply->error();
		const QString errorString = reply->errorString();
		const QByteArray data = reply->readAll();
		reply->deleteLater();

		int size = 0;
		if (error == QNetworkReply::NoError) {
			setServerError(QString());
			const QJsonDocument doc = QJsonDocument::fromJson(data);
			if (doc.isArray()) {
				for (const auto& value : doc.array()) {
					const QJsonObject slot = value.toObject();
					if (slot.contains("n_ctx")) {
						size = slot["n_ctx"].toInt(0);
						if (size > 0) break;
					}
				}
			} else if (doc.isObject()) {
				const QJsonObject obj = doc.object();
				size = obj["n_ctx"].toInt(0);
				if (size <= 0)
					size = obj["default_generation_settings"]
							   .toObject()["n_ctx"]
							   .toInt(0);
			}
		} else {
			const QString serverMessage = jsonErrorMessage(data);
			setServerError(
				serverMessage.isEmpty()
					? QStringLiteral("Cannot reach LLM server: %1")
						  .arg(errorString)
					: serverMessage);
		}
		if (size > 0)
			setContextSize(size);
		else if (m_contextSize <= 0)
			setContextSize(4096);
		if (done) done();
	};
	connect(reply, &QNetworkReply::finished, this, finish);
	QTimer::singleShot(kProbeTimeoutMs, this, [reply, settled]() {
		if (!*settled) reply->abort();
	});
}

void LlmClient::updateTokenUsage(const QJsonObject& data) {
	ChatSession* active = m_active;
	if (!active || m_contextSize <= 0) return;
	const QJsonObject usage = data["usage"].toObject();
	if (usage.isEmpty()) return;
	const double used = usage.value("prompt_tokens").toDouble() +
						usage.value("completion_tokens").toDouble();
	if (used > 0) active->setLastTokenCount(static_cast<int>(used));
}

void LlmClient::shortRequest(
	const QString& tag,
	const QString& systemPrompt,
	const QString& userText,
	std::function<void(QString result)> onResult) {
	const QUrl url =
		QUrl::fromUserInput(completionsPath(m_endpoint, "/chat/completions"));
	if (!url.isValid() || url.host().isEmpty()) return;

	qInfo() << "LlmClient:" << tag << "request POST" << url.toString()
			<< "model=" << m_model;

	QNetworkRequest request(url);
	request.setHeader(QNetworkRequest::ContentTypeHeader, "application/json");
	request.setRawHeader("Accept", "application/json");

	QJsonArray messages;
	QJsonObject system;
	system[QStringLiteral("role")] = QStringLiteral("system");
	system[QStringLiteral("content")] = systemPrompt;
	messages.append(system);
	QJsonObject user;
	user[QStringLiteral("role")] = QStringLiteral("user");
	user[QStringLiteral("content")] = userText.simplified().mid(0, 512);
	messages.append(user);

	QJsonObject body;
	if (!m_model.isEmpty()) body[QStringLiteral("model")] = m_model;
	body[QStringLiteral("stream")] = false;
	body[QStringLiteral("temperature")] = 0.3;
	body[QStringLiteral("max_tokens")] = 128;
	QJsonObject templateKwargs;
	templateKwargs[QStringLiteral("enable_thinking")] = false;
	body[QStringLiteral("chat_template_kwargs")] = templateKwargs;
	body[QStringLiteral("messages")] = messages;

	auto* reply = m_manager.post(request, QJsonDocument(body).toJson());
	connect(
		reply,
		&QNetworkReply::finished,
		this,
		[this, reply, tag, onResult = std::move(onResult)]() {
			const QByteArray data = reply->readAll();
			reply->deleteLater();

			qInfo() << "LlmClient:" << tag << "request finished"
					<< "error=" << reply->error() << reply->errorString()
					<< "http="
					<< reply
						   ->attribute(QNetworkRequest::HttpStatusCodeAttribute)
						   .toInt()
					<< "response="
					<< QString::fromUtf8(data.left(400)).simplified();

			if (reply->error() != QNetworkReply::NoError) return;

			const QJsonDocument doc = QJsonDocument::fromJson(data);
			const QJsonArray choices =
				doc.object()[QStringLiteral("choices")].toArray();
			if (choices.isEmpty()) return;
			const QString result = choices.at(0)
									   .toObject()[QStringLiteral("message")]
									   .toObject()[QStringLiteral("content")]
									   .toString()
									   .trimmed();
			qInfo() << "LlmClient:" << tag << "raw result" << result;
			onResult(result);
		});
}

void LlmClient::requestTitle(ChatSession* session, const QString& userText) {
	shortRequest(
		QStringLiteral("title"),
		QStringLiteral(
			"Write a short, concise title for a chat conversation starting "
			"with the user's message below. At most six words, no quotation "
			"marks, no trailing punctuation. Reply with the title only."),
		userText,
		[this, session = QPointer<ChatSession>(session)](QString title) {
			if (!session) {
				qWarning() << "LlmClient: title request: session gone";
				return;
			}
			const auto isQuote = [](QChar c) {
				return c == QLatin1Char('"') || c == QLatin1Char('\'') ||
					   c == QChar(u'\u201C') || c == QChar(u'\u201D') ||
					   c == QChar(u'\u2018') || c == QChar(u'\u2019');
			};
			while (title.size() >= 2 && isQuote(title.at(0)) &&
				   isQuote(title.at(title.size() - 1)))
				title = title.mid(1, title.size() - 2).simplified();
			while (!title.isEmpty() && (title.endsWith(QLatin1Char('.')) ||
										title.endsWith(QLatin1Char('!')) ||
										title.endsWith(QLatin1Char('?'))))
				title.chop(1);
			if (title.size() < 2) {
				qWarning() << "LlmClient: title rejected (too short)" << title;
				return;
			}
			if (title.size() > 48) title = title.left(47) + QStringLiteral("…");
			qInfo() << "LlmClient: suggesting title" << title;
			Q_EMIT titleSuggested(session, title);
		});
}

void LlmClient::requestIcon(ChatSession* session, const QString& userText) {
	const QStringList icons = {
		"chat",			  "lightbulb", "code",		   "description",
		"article",		  "school",	   "work",		   "build",
		"science",		  "palette",   "music_note",   "sports_esports",
		"takeout_dining", "flight",	   "photo_camera", "psychology_alt",
		"favorite",		  "savings",   "gamepad",	   "auto_awesome",
	};
	const QString prompt =
		QStringLiteral(
			"Pick the single icon name from this list that best matches the "
			"topic of the user's message below: %1. Reply with only the icon "
			"name, exactly as written in the list, and nothing else.")
			.arg(icons.join(QStringLiteral(", ")));
	shortRequest(
		QStringLiteral("icon"),
		prompt,
		userText,
		[this, session = QPointer<ChatSession>(session), icons](QString name) {
			if (!session) {
				qWarning() << "LlmClient: icon request: session gone";
				return;
			}
			name = name.simplified().toLower();
			const auto isQuote = [](QChar c) {
				return c == QLatin1Char('"') || c == QLatin1Char('\'');
			};
			while (name.size() >= 2 && isQuote(name.at(0)) &&
				   isQuote(name.at(name.size() - 1)))
				name = name.mid(1, name.size() - 2).simplified();
			name.replace(QLatin1Char(' '), QLatin1Char('_'));
			if (!icons.contains(name)) {
				qWarning() << "LlmClient: icon not in list, using default"
						   << name;
				name = QStringLiteral("chat");
			}
			qInfo() << "LlmClient: suggesting icon" << name;
			Q_EMIT iconSuggested(session, name);
		});
}

} // namespace ZShell::llm