#include "shelltool.hpp"

#include "llmclient.hpp"
#include "session.hpp"

#include <QDir>
#include <QJsonArray>
#include <QJsonObject>
#include <QProcess>
#include <QTimer>

#include <unistd.h>

namespace ZShell::llm {

namespace {

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

ShellTool::ShellTool(LlmClient* client, QObject* parent)
	: LlmTool(parent), m_client(client) {}

ShellTool::~ShellTool() {
	for (auto* job : m_jobs) {
		if (job->timer) job->timer->stop();
		if (job->process) job->process->kill();
		delete job;
	}
	m_jobs.clear();
}

QString ShellTool::name() const {
	return QStringLiteral("runcommand");
}

QString ShellTool::description() const {
	return QStringLiteral(
		"Run a shell command inside a bubblewrap sandbox that contains "
		"the files the user attached to the conversation. The rest of "
		"the filesystem is mounted read-only, there is no network "
		"access, and only the sandbox directory is writable. The "
		"command is executed with /bin/sh and the sandbox directory as "
		"its working directory, so relative paths refer to the sandbox. "
		"Use it to explore, transform, or extract attached files (for "
		"example: ls, find, grep, unzip, tar, head, wc). The combined "
		"output is limited; the exit code is always reported. The call "
		"only completes when every child process has exited, so do not "
		"start long-lived background processes.");
}

QJsonObject ShellTool::parameters() const {
	QJsonObject command;
	command[QStringLiteral("type")] = QStringLiteral("string");
	command[QStringLiteral("description")] =
		QStringLiteral("The shell command to execute");

	QJsonObject timeout;
	timeout[QStringLiteral("type")] = QStringLiteral("integer");
	timeout[QStringLiteral("minimum")] = 1;
	timeout[QStringLiteral("maximum")] = MaxTimeoutSeconds;
	timeout[QStringLiteral("description")] =
		QStringLiteral("Timeout in seconds. Defaults to %1.")
			.arg(DefaultTimeoutSeconds);

	QJsonObject properties;
	properties[QStringLiteral("command")] = command;
	properties[QStringLiteral("timeout")] = timeout;

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

QStringList ShellTool::sandboxArguments(const QString& workspace) {
	QStringList args;
	args << QLatin1String("--unshare-user") << QLatin1String("--uid")
		 << QString::number(::getuid()) << QLatin1String("--gid")
		 << QString::number(::getgid()) << QLatin1String("--disable-userns")
		 << QLatin1String("--unshare-pid") << QLatin1String("--unshare-net")
		 << QLatin1String("--unshare-ipc") << QLatin1String("--unshare-uts")
		 << QLatin1String("--hostname") << QLatin1String("zshell-sandbox");
	args << QLatin1String("--ro-bind") << QLatin1String("/")
		 << QLatin1String("/") << QLatin1String("--tmpfs")
		 << QLatin1String("/tmp") << QLatin1String("--tmpfs")
		 << QLatin1String("/run") << QLatin1String("--bind") << workspace
		 << workspace << QLatin1String("--dev") << QLatin1String("/dev")
		 << QLatin1String("--proc") << QLatin1String("/proc");
	args << QLatin1String("--chdir") << workspace;
	args << QLatin1String("--clearenv") << QLatin1String("--setenv")
		 << QLatin1String("PATH")
		 << QLatin1String(
				"/usr/local/sbin:/usr/local/bin:/usr/sbin:"
				"/usr/bin:/sbin:/bin")
		 << QLatin1String("--setenv") << QLatin1String("HOME") << workspace
		 << QLatin1String("--setenv") << QLatin1String("TMPDIR")
		 << QLatin1String("/tmp") << QLatin1String("--setenv")
		 << QLatin1String("SHELL") << QLatin1String("/bin/sh");
	const QByteArray user = qgetenv("USER");
	if (!user.isEmpty())
		args << QLatin1String("--setenv") << QLatin1String("USER")
			 << QString::fromLocal8Bit(user) << QLatin1String("--setenv")
			 << QLatin1String("LOGNAME") << QString::fromLocal8Bit(user);
	const QByteArray lang = qgetenv("LANG");
	if (!lang.isEmpty())
		args << QLatin1String("--setenv") << QLatin1String("LANG")
			 << QString::fromLocal8Bit(lang);
	args << QLatin1String("--new-session")
		 << QLatin1String("--die-with-parent");
	return args;
}

void ShellTool::execute(
	const QString& toolCallId,
	const QJsonObject& args,
	int outputBudgetChars,
	std::function<void(const QJsonObject&)> done) {
	Q_UNUSED(toolCallId);
	const QString command = args[QStringLiteral("command")].toString();
	if (command.trimmed().isEmpty()) {
		done(makeError(QStringLiteral("Missing required argument 'command'")));
		return;
	}
	LlmClient* client = m_client;
	if (!client || !client->sandboxEnabled()) {
		done(makeError(QStringLiteral(
			"The file sandbox is disabled, so no commands can run")));
		return;
	}
	const QString workspace = client->activeWorkspace();
	if (workspace.isEmpty()) {
		done(makeError(
			QStringLiteral("No sandbox available: no chat session is active")));
		return;
	}
	QDir().mkpath(workspace);

	const int timeoutSeconds = qBound(
		1,
		args[QStringLiteral("timeout")].toInt(DefaultTimeoutSeconds),
		MaxTimeoutSeconds);
	auto* job = new Job;
	job->timeoutSeconds = timeoutSeconds;
	job->outputBudgetChars = outputBudgetChars;
	job->done = std::move(done);
	m_jobs.append(job);

	QProcess* process = new QProcess(this);
	process->setWorkingDirectory(workspace);
	process->setProcessChannelMode(QProcess::MergedChannels);
	job->process = process;

	QTimer* timer = new QTimer(this);
	timer->setSingleShot(true);
	connect(timer, &QTimer::timeout, this, [this, job, process]() {
		job->timedOut = true;
		process->kill();
	});
	timer->start(timeoutSeconds * 1000);
	job->timer = timer;

	connect(
		process,
		QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished),
		this,
		[this, job](int, QProcess::ExitStatus) { finishJob(job); });
	connect(
		process,
		&QProcess::errorOccurred,
		this,
		[this, job, process](QProcess::ProcessError error) {
			if (job->finished) return;
			if (error == QProcess::FailedToStart ||
				process->state() == QProcess::NotRunning) {
				job->startError = process->errorString();
				finishJob(job);
			} else {
				process->kill();
			}
		});

	QStringList commandArgs = sandboxArguments(workspace);
	commandArgs << QStringLiteral("sh") << QStringLiteral("-c") << command;
	process->start(QStringLiteral("bwrap"), commandArgs);
}

void ShellTool::cancel() {
	for (auto* job : m_jobs) {
		if (job->finished || !job->process) continue;
		job->cancelled = true;
		job->process->kill();
	}
}

void ShellTool::finishJob(Job* job) {
	if (!job || job->finished) return;
	job->finished = true;
	if (job->timer) job->timer->stop();
	m_jobs.removeAll(job);

	QProcess* process = job->process;
	const bool ran = process && process->state() != QProcess::NotRunning;
	const int exitCode = ran ? process->exitCode() : -1;
	QString output;
	if (process) output = QString::fromUtf8(process->readAllStandardOutput());

	QString note;
	if (!job->startError.isEmpty())
		note =
			QStringLiteral("Command failed to start: %1\n").arg(job->startError);
	else if (job->cancelled)
		note = QStringLiteral("Command was cancelled.\n");
	else if (job->timedOut)
		note = QStringLiteral(
				   "Command timed out after %1 seconds and was killed.\n")
				   .arg(job->timeoutSeconds);

	int budget = inlineBudgetChars();
	if (job->outputBudgetChars > 0)
		budget = qMin(budget, job->outputBudgetChars);
	if (output.size() > budget)
		output = output.left(budget) + QStringLiteral(
										   "\n[output truncated: %1 characters "
										   "total]")
										   .arg(output.size());

	const QString result =
		note + QStringLiteral("exit code: %1\n").arg(exitCode) +
		(output.isEmpty() ? QStringLiteral("(no output)") : output);

	QJsonObject resultObject;
	if (!job->cancelled && !job->timedOut && exitCode == 0)
		resultObject = makeOutput(result);
	else
		resultObject = makeError(result);

	std::function<void(const QJsonObject&)> done = std::move(job->done);
	if (process) process->deleteLater();
	delete job;
	if (done) done(resultObject);
}

} // namespace ZShell::llm