#include "exllama.h"

#include <QJsonObject>
#include <QJsonDocument>

ExLlama::ExLlama()
{
	connect(&m_webSocket, &QWebSocket::textMessageReceived, this, &ExLlama::socketMessage);
}

bool ExLlama::ready()
{
	return m_webSocket.isValid();
}

void ExLlama::socketMessage(const QString& message)
{
	QJsonDocument jsonDocument = QJsonDocument::fromJson(message.toUtf8());
	QJsonValue idVal = jsonDocument[QStringLiteral("request_id")];
	if(!idVal.isDouble())
	{
		qDebug()<<"Got invalid response on socket";
		return;
	}
	int id = idVal.toInt();

	if(!isValidId(id))
	{
		qDebug()<<"Got unkown response id on socket";
		return;
	}

	QJsonValue responseValue = jsonDocument[QStringLiteral("response")];
	if(!responseValue.isString())
	{
		qDebug()<<"Got invalid response on socket";
		return;
	}

	feedResponse(Response(responseValue.toString(), id, true));
}

void ExLlama::generateImpl(const Request& request)
{
	QJsonObject json;
	json[QStringLiteral("action")] = QStringLiteral("infer");
	json[QStringLiteral("request_id")] = static_cast<double>(request.getId());
	json[QStringLiteral("text")] = request.getText();
	json[QStringLiteral("max_new_tokens")] = 50;
	json[QStringLiteral("stream")] = false;

	QJsonDocument jsonDocument(json);
	QString requestText = QString::fromUtf8(jsonDocument.toJson(QJsonDocument::JsonFormat::Compact));
	qDebug()<<__func__<<' '<<requestText;
	m_webSocket.sendTextMessage(requestText);
}

void ExLlama::open(const QUrl& url)
{
	m_webSocket.close();
	m_webSocket.open(url);
}

ExLlama::~ExLlama()
{
	m_webSocket.disconnect();
}