#include "webfetchtool.hpp"

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

namespace ZShell::llm {

const QString WebFetchTool::StoragePath =
	QStringLiteral("/tmp/zshell-llm/webfetch");

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. Only a limited amount of content is returned inline; "
		"the full output of every call is saved under "
		"/tmp/zshell-llm/webfetch/, where it can be paged through with "
		"the readfile tool. 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 QString& toolCallId,
	const QJsonObject& args,
	std::function<void(const QJsonObject&)> done) {
	auto* job = new Job;
	job->done = std::move(done);
	job->toolCallId = toolCallId;
	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);

		const QString savedPath = saveToFile(*job, content, mime);

		QString output = content;
		const int budget = inlineBudgetChars();
		if (content.size() > budget) {
			output = content.left(budget);
			const QString note = QStringLiteral(
				"\n\n[... truncated: showing %1 of %2 characters")
								.arg(budget)
								.arg(content.size());
			if (!savedPath.isEmpty())
				output += note +
						  QStringLiteral(
							  ". The full content is saved to %1; use "
							  "the readfile tool to read the rest.")
							  .arg(savedPath);
			else
				output += note + QStringLiteral(
								   ". The remaining content is not "
								   "available.");
		}
		completeJob(job, makeOutput(output));
	});
}

QString WebFetchTool::saveToFile(
	const Job& job, const QString& content, const QString& mime) const {
	QDir dir(StoragePath);
	if (!dir.exists() && !dir.mkpath(QStringLiteral(".")))
		return QString();

	QString fileName;
	fileName.reserve(job.toolCallId.size());
	for (const QChar& c : job.toolCallId)
		fileName += (c.isLetterOrNumber() || c == QLatin1Char('-') ||
					  c == QLatin1Char('_'))
						? c
						: QLatin1Char('_');
	if (fileName.isEmpty()) fileName = QStringLiteral("fetch");

	QString extension = QStringLiteral("txt");
	if (job.format == QLatin1String("html"))
		extension = QStringLiteral("html");
	else if (mime.contains(QLatin1String("json")))
		extension = QStringLiteral("json");
	else if (mime.contains(QLatin1String("xml")))
		extension = QStringLiteral("xml");

	const QString path =
		dir.filePath(fileName + QLatin1Char('.') + extension);
	QFile file(path);
	if (!file.open(QIODevice::WriteOnly | QIODevice::Truncate))
		return QString();
	file.write(content.toUtf8());
	return path;
}

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
