#include "webfetchtool.hpp"

#include <QJsonArray>
#include <QNetworkReply>
#include <QNetworkRequest>
#include <QSet>
#include <QTimer>

namespace ZShell::llm {

namespace {

const char* kUserAgent =
	"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
	"(KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36";

QJsonObject makeOutput(const QString& text) {
	QJsonObject obj;
	obj[QStringLiteral("output")] = text;
	return obj;
}

QJsonObject makeError(const QString& message) {
	QJsonObject obj;
	obj[QStringLiteral("error")] = message;
	return obj;
}

} // namespace

WebFetchTool::WebFetchTool(QObject* parent) : LlmTool(parent) {}

WebFetchTool::~WebFetchTool() {
	for (auto* job : m_jobs) {
		if (job->timer) job->timer->stop();
		if (job->reply) job->reply->abort();
	}
	qDeleteAll(m_jobs);
}

QString WebFetchTool::name() const {
	return QStringLiteral("webfetch");
}

QString WebFetchTool::description() const {
	return QStringLiteral(
		"Fetch content from an HTTP or HTTPS URL and return it as plain "
		"text or raw HTML. HTML pages are reduced to their visible text "
		"by default. This tool is read-only.");
}

QJsonObject WebFetchTool::parameters() const {
	QJsonObject url;
	url[QStringLiteral("type")] = QStringLiteral("string");
	url[QStringLiteral("description")] =
		QStringLiteral("The HTTP or HTTPS URL to fetch content from");

	QJsonArray formats;
	formats.append(QStringLiteral("text"));
	formats.append(QStringLiteral("html"));
	QJsonObject format;
	format[QStringLiteral("type")] = QStringLiteral("string");
	format[QStringLiteral("enum")] = formats;
	format[QStringLiteral("description")] = QStringLiteral(
		"The format to return the content in. Defaults to "
		"text.");

	QJsonObject timeout;
	timeout[QStringLiteral("type")] = QStringLiteral("integer");
	timeout[QStringLiteral("minimum")] = 1;
	timeout[QStringLiteral("maximum")] = MaxTimeoutSeconds;
	timeout[QStringLiteral("description")] =
		QStringLiteral("Optional timeout in seconds");

	QJsonObject properties;
	properties[QStringLiteral("url")] = url;
	properties[QStringLiteral("format")] = format;
	properties[QStringLiteral("timeout")] = timeout;

	QJsonObject schema;
	schema[QStringLiteral("type")] = QStringLiteral("object");
	schema[QStringLiteral("properties")] = properties;
	QJsonArray required;
	required.append(QStringLiteral("url"));
	schema[QStringLiteral("required")] = required;
	return schema;
}

void WebFetchTool::completeJob(Job* job, QJsonObject result) {
	QMetaObject::invokeMethod(
		this,
		[this, job, result = std::move(result)]() mutable {
			if (!m_jobs.contains(job)) return;
			auto done = std::move(job->done);
			m_jobs.removeAll(job);
			job->timer->deleteLater();
			if (job->reply) job->reply->deleteLater();
			delete job;
			done(std::move(result));
		},
		Qt::QueuedConnection);
}

void WebFetchTool::execute(
	const QJsonObject& args, std::function<void(const QJsonObject&)> done) {
	auto* job = new Job;
	job->done = std::move(done);
	m_jobs.append(job);

	auto fail = [this, job](const QString& message) {
		job->timer->stop();
		completeJob(job, makeError(message));
	};

	const QString urlText = args[QStringLiteral("url")].toString().trimmed();
	const QUrl url = QUrl::fromUserInput(urlText);
	if (!url.isValid() || url.host().isEmpty()) {
		fail(QStringLiteral("Invalid URL: %1").arg(urlText));
		return;
	}
	if (url.scheme() != QLatin1String("http") &&
		url.scheme() != QLatin1String("https")) {
		fail(QStringLiteral("URL must use http:// or https://"));
		return;
	}

	job->format =
		args[QStringLiteral("format")].toString(QStringLiteral("text"));
	if (job->format != QLatin1String("html"))
		job->format = QStringLiteral("text");

	const int timeoutMs =
		qBound(
			1,
			args[QStringLiteral("timeout")].toInt(DefaultTimeoutSeconds),
			MaxTimeoutSeconds) *
		1000;

	job->timer = new QTimer(this);
	job->timer->setSingleShot(true);
	connect(job->timer, &QTimer::timeout, this, [this, job]() {
		if (job->reply) job->reply->abort();
	});
	job->timer->start(timeoutMs);

	QNetworkRequest request(url);
	request.setRawHeader("User-Agent", QByteArray(kUserAgent));
	request.setRawHeader(
		"Accept",
		job->format == QLatin1String("html")
			? "text/html;q=1.0, application/xhtml+xml;q=0.9, */*;q=0.1"
			: "text/plain;q=1.0, text/markdown;q=0.9, text/html;q=0.8, "
			  "*/*;q=0.1");
	request.setRawHeader("Accept-Language", "en-US,en;q=0.9");

	job->reply = m_manager.get(request);
	connect(job->reply, &QNetworkReply::readyRead, this, [this, job]() {
		if (!job->reply) return;
		job->body += job->reply->readAll();
		if (job->body.size() > MaxResponseBytes) {
			job->tooLarge = true;
			job->reply->abort();
		}
	});
	connect(job->reply, &QNetworkReply::finished, this, [this, job]() {
		QNetworkReply* reply = job->reply;
		job->reply = nullptr;
		job->timer->stop();
		if (!reply) return;

		const QNetworkReply::NetworkError error = reply->error();
		const QString errorString = reply->errorString();
		QByteArray body = job->body;
		body += reply->readAll();
		job->body.clear();
		const QByteArray contentType =
			reply->rawHeader("Content-Type").toLower();
		const int status =
			reply->attribute(QNetworkRequest::HttpStatusCodeAttribute).toInt();

		if (job->tooLarge) {
			completeJob(
				job,
				makeError(QStringLiteral(
					"Response too large (exceeds the 5 MB "
					"limit")));
			return;
		}
		if (error != QNetworkReply::NoError) {
			completeJob(
				job,
				makeError(
					QStringLiteral("Request failed: %1").arg(errorString)));
			return;
		}
		if (status < 200 || status >= 300) {
			completeJob(
				job,
				makeError(
					QStringLiteral("Server returned status %1").arg(status)));
			return;
		}

		const QString mime = QString::fromLatin1(contentType)
								 .section(QLatin1Char(';'), 0, 0)
								 .trimmed();
		if (mime.startsWith(QLatin1String("image/"))) {
			completeJob(
				job,
				makeError(
					QStringLiteral("Unsupported fetched image content type: %1")
						.arg(mime)));
			return;
		}
		const bool textual =
			mime.isEmpty() || mime.startsWith(QLatin1String("text/")) ||
			mime == QLatin1String("application/json") ||
			mime.endsWith(QLatin1String("+json")) ||
			mime == QLatin1String("application/xml") ||
			mime.endsWith(QLatin1String("+xml")) ||
			mime.startsWith(QLatin1String("application/javascript")) ||
			mime.startsWith(QLatin1String("text/javascript"));
		if (!textual) {
			completeJob(
				job,
				makeError(
					QStringLiteral("Unsupported fetched file content type: %1")
						.arg(mime)));
			return;
		}

		QString content = QString::fromUtf8(body);
		if (mime.contains(QLatin1String("text/html")) &&
			job->format == QLatin1String("text"))
			content = extractTextFromHtml(content);
		if (content.size() > MaxOutputChars)
			content = content.left(MaxOutputChars) +
					  QStringLiteral("\n[... truncated ...]");
		completeJob(job, makeOutput(content));
	});
}

void WebFetchTool::cancel() {
	for (auto* job : m_jobs) {
		job->timer->stop();
		if (job->reply) job->reply->abort();
	}
}

QString WebFetchTool::decodeEntities(const QString& text) {
	if (!text.contains(QLatin1Char('&'))) return text;
	QString out;
	out.reserve(text.size());
	for (qsizetype i = 0; i < text.size(); ++i) {
		if (text.at(i) != QLatin1Char('&')) {
			out += text.at(i);
			continue;
		}
		const qsizetype semi = text.indexOf(QLatin1Char(';'), i);
		if (semi < 0 || semi - i > 12) {
			out += QLatin1Char('&');
			continue;
		}
		const QString entity = text.mid(i + 1, semi - i - 1);
		QString replacement;
		if (entity == QLatin1String("amp"))
			replacement = QLatin1Char('&');
		else if (entity == QLatin1String("lt"))
			replacement = QLatin1Char('<');
		else if (entity == QLatin1String("gt"))
			replacement = QLatin1Char('>');
		else if (entity == QLatin1String("quot"))
			replacement = QLatin1Char('"');
		else if (entity == QLatin1String("apos"))
			replacement = QLatin1Char('\'');
		else if (entity == QLatin1String("nbsp"))
			replacement = QLatin1Char(' ');
		else {
			bool ok = false;
			const quint32 codePoint =
				entity.startsWith(QLatin1String("#x")) ||
						entity.startsWith(QLatin1String("#X"))
					? entity.mid(2).toUInt(&ok, 16)
					: entity.toUInt(&ok);
			if (ok && codePoint != 0) {
				const char32_t ucs4[2] = {
					static_cast<char32_t>(codePoint),
					0,
				};
				replacement = QString::fromUcs4(ucs4);
			}
		}
		if (replacement.isEmpty()) {
			out += QLatin1Char('&');
			continue;
		}
		out += replacement;
		i = semi;
	}
	return out;
}

QString WebFetchTool::extractTextFromHtml(const QString& html) {
	static const QSet<QString> kSkipTags = {
		QStringLiteral("noscript"),
		QStringLiteral("iframe"),
		QStringLiteral("object"),
		QStringLiteral("head"),
	};
	static const QSet<QString> kRawTags = {
		QStringLiteral("script"),
		QStringLiteral("style"),
	};
	static const QSet<QString> kVoidTags = {
		QStringLiteral("area"),
		QStringLiteral("base"),
		QStringLiteral("br"),
		QStringLiteral("col"),
		QStringLiteral("embed"),
		QStringLiteral("hr"),
		QStringLiteral("img"),
		QStringLiteral("input"),
		QStringLiteral("link"),
		QStringLiteral("meta"),
		QStringLiteral("source"),
		QStringLiteral("track"),
		QStringLiteral("wbr"),
	};

	QString text;
	text.reserve(html.size() / 2);
	int skipDepth = 0;
	qsizetype i = 0;
	while (i < html.size()) {
		const qsizetype open = html.indexOf(QLatin1Char('<'), i);
		if (open < 0) {
			if (skipDepth == 0) text += html.mid(i);
			break;
		}
		if (skipDepth == 0) text += html.mid(i, open - i);
		const qsizetype close = html.indexOf(QLatin1Char('>'), open);
		if (close < 0) break;
		const QString tag =
			html.mid(open + 1, close - open - 1).trimmed().toLower();
		i = close + 1;

		if (tag.startsWith(QLatin1Char('!')) ||
			tag.startsWith(QLatin1Char('?')))
			continue;

		QString name = tag;
		if (name.startsWith(QLatin1Char('/'))) {
			if (skipDepth > 0) --skipDepth;
			continue;
		}
		qsizetype j = 0;
		while (j < name.size() && (name.at(j).isLetterOrNumber() ||
								   name.at(j) == QLatin1Char(':') ||
								   name.at(j) == QLatin1Char('-')))
			++j;
		name = name.left(j);

		if (kRawTags.contains(name)) {
			const qsizetype rawEnd = html.indexOf(
				QStringLiteral("</") + name, i, Qt::CaseInsensitive);
			if (rawEnd < 0) break;
			const qsizetype rawClose = html.indexOf(QLatin1Char('>'), rawEnd);
			if (rawClose < 0) break;
			i = rawClose + 1;
			continue;
		}
		if (kVoidTags.contains(name)) continue;
		if (skipDepth > 0) {
			if (name == QLatin1String("body")) {
				skipDepth = 0;
				continue;
			}
			++skipDepth;
			continue;
		}
		if (kSkipTags.contains(name)) {
			++skipDepth;
			continue;
		}
		text += QLatin1Char(' ');
	}

	QString out = decodeEntities(text);
	QStringList lines;
	for (const QString& line : out.split(QLatin1Char('\n'))) {
		const QString flat = line.simplified();
		if (flat.isEmpty()) {
			if (!lines.isEmpty() && lines.last().isEmpty()) continue;
			lines.append(QString());
		} else {
			lines.append(flat);
		}
	}
	while (lines.size() > 1 && lines.first().isEmpty())
		lines.removeFirst();
	while (lines.size() > 1 && lines.last().isEmpty())
		lines.removeLast();
	return lines.join(QLatin1Char('\n'));
}

} // namespace ZShell::llm