#include "kateai.h"

#include <QShortcut>
#include <KConfigGroup>
#include <KPluginFactory>
#include <KTextEditor/Document>
#include <KTextEditor/View>
#include <qdebug.h>
#include <qhash.h>
#include <qjsonobject.h>
#include <qnamespace.h>
#include <QString>
#include <KActionCollection>
#include <KXMLGUIFactory>
#include <QJsonObject>
#include <QJsonDocument>
#include <QRandomGenerator>
#include <QTextCodec>
#include <limits>
#include <QMessageBox>
#include <KLocalizedString>
#include <KSharedConfig>

#include "kateaiconfigpage.h"

K_PLUGIN_FACTORY_WITH_JSON(KateAiPluginFactory, "kateai.json", registerPlugin<KateAiPlugin>();)

KateAiPlugin::KateAiPlugin(QObject *parent, const QList<QVariant> &)
	: KTextEditor::Plugin(parent), m_serverUrl(QStringLiteral("ws://localhost:8642"))
{
	connect(&m_webSocket, &QWebSocket::connected, this, &KateAiPlugin::onConnected);
	readConfig();
}

void KateAiPlugin::onConnected()
{
	qDebug()<<__func__<<m_webSocket.isValid();
}

KateAiPlugin::~KateAiPlugin() = default;

void KateAiPlugin::reconnect()
{
	m_webSocket.close();
	m_webSocket.open(m_serverUrl);
}

QObject *KateAiPlugin::createView(KTextEditor::MainWindow *mainWindow)
{
	auto view = new KateAiPluginView(this, mainWindow, &m_webSocket);
	connect(view, &KateAiPluginView::reconnect, this, &KateAiPlugin::reconnect);
	return view;
}

KTextEditor::ConfigPage *KateAiPlugin::configPage(int number, QWidget *parent)
{
	if (number != 0)
		return nullptr;

	return new KateAiConfigPage(parent, this);
}

void KateAiPlugin::readConfig()
{
	KConfigGroup config(KSharedConfig::openConfig(), "Ai");
	m_serverUrl = QUrl(config.readEntry("Url", "ws://localhost:8642"));
	reconnect();
}

KateAiPluginView::KateAiPluginView(KateAiPlugin *plugin, KTextEditor::MainWindow *mainwindow,
                                   QPointer<QWebSocket> webSocket, bool instruct)
	: QObject(plugin)
	, m_mainWindow(mainwindow)
	, m_webSocket(webSocket)
	, m_useInstruct(instruct)
{
	KXMLGUIClient::setComponentName(QStringLiteral("kateaiplugin"), QStringLiteral("Git Blame"));
	setXMLFile(QStringLiteral("ui.rc"));
	QAction *generateAction = actionCollection()->addAction(QStringLiteral("ai_generate"));
	generateAction->setText(QStringLiteral("Generate text using AI"));
	actionCollection()->setDefaultShortcut(generateAction, QKeySequence(Qt::CTRL | Qt::ALT | Qt::Key_A));
	m_mainWindow->guiFactory()->addClient(this);

	connect(generateAction, &QAction::triggered, this, &KateAiPluginView::generate);
	connect(m_webSocket, &QWebSocket::textMessageReceived, this, &KateAiPluginView::socketMessage);
}

QPointer<KTextEditor::Document> KateAiPluginView::activeDocument() const
{
	KTextEditor::View *view = m_mainWindow->activeView();
	if(view && view->document())
		return view->document();
	return nullptr;
}

KTextEditor::Cursor KateAiPluginView::getCurrentCursor() const
{
	KTextEditor::View *view = m_mainWindow->activeView();
	if(view)
		return view->cursorPosition();
	return KTextEditor::Cursor();
}

void KateAiPluginView::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();

	QHash<int, Request>::iterator it = m_requests.find(id);
	if(it != m_requests.end())
	{
		QJsonValue responseValue = jsonDocument[QStringLiteral("response")];
		if(!responseValue.isString())
		{
			qDebug()<<"Got invalid response on socket";
			return;
		}
		if(it.value().document)
			it.value().document->insertText(it.value().cursor, responseValue.toString());
		m_requests.erase(it);
	}
}

QStringList KateAiPluginView::getIncludePaths(const QString& text)
{
	QStringList lines = text.split(U'\n');
	QStringList paths;

	for(const QString& line : lines)
	{
		if(line.trimmed().startsWith(QStringLiteral("#include")))
		{
			int start = line.indexOf(U'<');
			int end;
			if(start != -1)
			{
				end = line.indexOf(U'>', start+1);
				if(end == -1)
					continue;
			}
			else
			{
				start = line.indexOf(U'"');
				if(start == -1)
					continue;

				end = line.indexOf(U'"', start+1);
				if(end == -1)
					continue;
			}
			paths.push_back(line.mid(start+1, (end-(start+1))).trimmed());
		}
	}
	return paths;
}

QString KateAiPluginView::assembleContext(QPointer<KTextEditor::Document> document, const KTextEditor::Cursor& cursor)
{
	QString mime = document->mimeType();
	QString context;
	QString baseText;

	if(!m_useInstruct)
		baseText = document->text(KTextEditor::Range(KTextEditor::Cursor(0, 0), cursor));
	else
		baseText = document->text();
	if(mime ==  QStringLiteral("text/x-c++src") || mime ==  QStringLiteral("text/x-csrc"))
	{
		QFileInfo documentFileInfo(document->url().path());
		QString directory = documentFileInfo.absolutePath();
		QStringList paths = getIncludePaths(baseText);
		qDebug()<<__func__<<"Directory:"<<directory<<"Paths:"<<paths;

		for(QString& path : paths)
		{
			path = directory + QDir::separator() + path;
			qDebug()<<path;
			QFile file(path);
			if(!file.isOpen())
				continue;
			QByteArray fileData = file.readAll();
			QString fileText = QString::fromUtf8(fileData);
			context.append(fileText);
		}
	}

	context.append(baseText);

	return context;
}

void KateAiPluginView::generate()
{
	qDebug()<<activeDocument()->mimeType();
	if(m_webSocket && m_webSocket->isValid())
	{
		KTextEditor::Cursor cursor = getCurrentCursor();
		QPointer<KTextEditor::Document> document = activeDocument();
		QString text = assembleContext(document, cursor);
		int id = QRandomGenerator::global()->bounded(0, std::numeric_limits<int>::max());

		QJsonObject json;
		json[QStringLiteral("action")] = QStringLiteral("infer");
		json[QStringLiteral("request_id")] = id;
		json[QStringLiteral("text")] = text;
		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);
		m_requests.insert(id, {cursor, document});
	}
	else
	{
		QMessageBox box;
		box.setText(i18n("The AI server is not connected."));
		box.setInformativeText(i18n("would you like to try and reconnect?"));
		box.setStandardButtons(QMessageBox::Yes | QMessageBox::No);
		box.setDefaultButton(QMessageBox::Yes);
		int ret = box.exec();
		if(ret == QMessageBox::Yes)
			reconnect();
	}
}

KateAiPluginView::~KateAiPluginView()
{
	m_mainWindow->guiFactory()->removeClient(this);
}

void KateAiPluginView::setInstruct(bool instruct)
{
	m_useInstruct = instruct;
}

#include "kateai.moc"
#include "moc_kateai.cpp"