z-bar-qt/Plugins/ZShell/Llm/llmclient.hpp

#pragma once

#include "segment.hpp"
#include "tool.hpp"

#include <QByteArray>
#include <QJsonArray>
#include <QNetworkAccessManager>
#include <QObject>
#include <QPointer>
#include <QString>
#include <QStringList>
#include <QTimer>

#include <functional>

class QJsonObject;
class QNetworkReply;

namespace ZShell::llm {

class ChatGeneration;
class ChatSession;
class LlmSegment;
class LlmTool;

class LlmClient : public QObject {
	Q_OBJECT

	public:
	explicit LlmClient(QObject* parent = nullptr);
	~LlmClient() override;

	[[nodiscard]] QString endpoint() const { return m_endpoint; }
	[[nodiscard]] QString model() const { return m_model; }
	[[nodiscard]] double temperature() const { return m_temperature; }
	[[nodiscard]] QStringList availableModels() const {
		return m_availableModels;
	}
	[[nodiscard]] int contextSize() const { return m_contextSize; }
	[[nodiscard]] ToolRegistry* tools() const { return m_tools; }
	void setEndpoint(const QString& value);
	void setModel(const QString& value);
	void setTemperature(double value);
	void setContextSize(int size);
	void probeContextSize(std::function<void()> done = {});
	Q_INVOKABLE void refreshFromServer();

	[[nodiscard]] bool busy() const { return m_busy; }
	[[nodiscard]] bool toolsEnabled() const { return m_tools->enabled(); }
	void setToolsEnabled(bool value) { m_tools->setEnabled(value); }
	[[nodiscard]] QString streamingChatId() const { return m_streamingChatId; }
	[[nodiscard]] ChatSession* streamingSession() const { return m_active; }

	void startGeneration(ChatSession* session, ChatGeneration* target);
	void stop();
	void endStream();
	void clearOnFinish(ChatSession* session);
	void sessionRemoved(ChatSession* session);

	void approveTools();
	void denyTools();

	void refreshModels();
	void requestTitle(ChatSession* session, const QString& userText);
	void requestIcon(ChatSession* session, const QString& userText);

	Q_SIGNALS:
	void busyChanged();
	void endpointChanged();
	void modelChanged();
	void availableModelsChanged();
	void contextSizeChanged();
	void toolsEnabledChanged();
	void streamingChatIdChanged();
	void errorOccurred(const QString& message);
	void titleSuggested(ZShell::llm::ChatSession* session, const QString& title);
	void iconSuggested(ZShell::llm::ChatSession* session, const QString& icon);

	private:
	struct ToolCallBuilder {
		QString id;
		QString name;
		QString arguments;
		QPointer<LlmSegment> segment;
		bool seen = false;
	};

	void sendRound();
	QJsonArray buildContextMessages(
		ChatSession* session, int stopBeforeRow, int extraReserveTokens) const;
	void applyToolCallDelta(const QJsonObject& call);
	void roundFinished();
	void executeAllCalls();
	void flushCallResults();
	void setApprovalPending(bool value);
	void finishPendingCalls(const QString& resultText);
	void finishTurn();
	void fail(const QString& message);
	void handleLine(const QByteArray& line);
	void drainBuffer();
	void updateTokenUsage(const QJsonObject& data);
	void setBusy(bool value);
	void setStreamingChatId(const QString& id);
	void shortRequest(
		const QString& tag,
		const QString& systemPrompt,
		const QString& userText,
		std::function<void(QString result)> onResult);
	static QString completionsPath(
		const QString& endpoint, const QString& subpath);
	static QString serverErrorMessage(
		const QByteArray& body, const QString& fallback);

	QNetworkAccessManager m_manager;
	ToolRegistry* m_tools = nullptr;
	QNetworkReply* m_reply = nullptr;
	QByteArray m_buffer;
	QTimer m_refreshTimer;
	QPointer<ChatSession> m_active;
	QPointer<ChatGeneration> m_streaming;
	QPointer<ChatSession> m_pendingClear;
	bool m_busy = false;
	QString m_streamingChatId;
	QString m_endpoint;
	QString m_model;
	QStringList m_availableModels;
	double m_temperature = 0.7;
	int m_contextSize = 0;

	QJsonArray m_transcript;
	QList<ToolCallBuilder> m_callBuilders;
	struct ToolCallResult {
		QString content;
		bool success = false;
	};
	QList<ToolCallResult> m_callResults;
	QString m_finishReason;
	int m_round = 0;
	qsizetype m_contentMark = 0;
	qsizetype m_reasoningMark = 0;
	bool m_roundDone = false;
	bool m_toolPhase = false;
	bool m_approvalPending = false;
	bool m_contextRetryUsed = false;
	int m_pendingCalls = 0;
	static constexpr int kMaxToolRounds = 12;
};

} // namespace ZShell::llm