// Copyright Microsoft and CHERIoT Contributors.
// SPDX-License-Identifier: MIT

#include <compartment.h>
#include <debug.hh>
#include <fail-simulator-on-error.h>
#include <microvium/microvium.h>

/// Expose debugging features unconditionally for this compartment.
using Debug = ConditionalDebug<true, "JavaScript hello compartment">;

namespace
{
	/**
	 * JavaScript bytecode.  The initialiser is generated by running
	 *
	 * ```
	 * microvium --output-bytes hello.js > bytecode.inc
	 * ```
	 */
	uint8_t bytecode[] =
#include "bytecode.inc"
	  ;

	/// Constant for the print function exposed to JavaScript->C++ FFI
	static constexpr mvm_HostFunctionID ImportPrint = 1;

	/// Constant for the say_hello function exposed to C++->JavaScript FFI
	static constexpr mvm_VMExportID ExportHello = 1234;

	/**
	 * Print a string passed from JavaScript.
	 */
	mvm_TeError print(mvm_VM            *vm,
	                  mvm_HostFunctionID funcID,
	                  mvm_Value         *result,
	                  mvm_Value         *args,
	                  uint8_t            argCount)
	{
		// Iterate over the arguments.
		for (unsigned i = 0; i < argCount; i++)
		{
			// Coerce the argument to a string and get it as a C string
			const char *str = mvm_toStringUtf8(vm, args[i], nullptr);
			// Write each character to the UART
			for (; *str != '\0'; str++)
			{
				ImplicitUARTOutput::write(*str);
			}
			// Try uncommenting the following line to see how it impacts total
			// heap usage.
			// mvm_runGC(vm, false);
		}
		// Write a trailing newline
		ImplicitUARTOutput::write('\n');
		// Unconditionally return success
		return MVM_E_SUCCESS;
	}

	/**
	 * Callback from microvium that resolves imports.  This just resolves the
	 * `print` function as the import corresponding to the `ImportPrint` ID.
	 */
	mvm_TeError resolve_import(mvm_HostFunctionID  funcID,
	                           void               *context,
	                           mvm_TfHostFunction *out)
	{
		if (funcID == ImportPrint)
		{
			*out = print;
			return MVM_E_SUCCESS;
		}
		return MVM_E_UNRESOLVED_IMPORT;
	}

	/**
	 * Helper that deletes a Microvium VM when used with a C++ unique pointer.
	 */
	struct MVMDeleter
	{
		void operator()(mvm_VM *mvm) const
		{
			mvm_free(mvm);
		}
	};
} // namespace

/// Thread entry point.
void __cheri_compartment("hello") say_hello()
{
	mvm_TeError                         err;
	std::unique_ptr<mvm_VM, MVMDeleter> vm;
	// Load the JavaScript bytecode snapshot
	{
		mvm_VM *rawVm;
		err = mvm_restore(&rawVm,
		                  bytecode,
		                  sizeof(bytecode),
		                  MALLOC_CAPABILITY,
		                  ::resolve_import);
		Debug::Assert(
		  err == MVM_E_SUCCESS, "Failed to parse bytecode: {}", err);
		vm.reset(rawVm);
	}

	// Get a handle to the JavaScript `say_hello` function.
	mvm_Value sayHello;
	err = mvm_resolveExports(vm.get(), &ExportHello, &sayHello, 1);
	Debug::Assert(
	  err == MVM_E_SUCCESS, "Failed to get say_hello function: {}", err);

	// Call the function:
	err = mvm_call(vm.get(), sayHello, nullptr, nullptr, 0);
	Debug::Assert(
	  err == MVM_E_SUCCESS, "Failed to call say_hello function: {}", err);

	mvm_TsMemoryStats stats;
	// Helper lambda to report the current memory usage.
	auto reportMemory = [&]() {
		mvm_getMemoryStats(vm.get(), &stats);
		Debug::log(
		  "Microvium is using {} bytes of memory, including {} bytes of heap",
		  stats.totalSize,
		  stats.virtualHeapUsed);
	};

	// See how much memory is used at the end of the invocation
	reportMemory();

	// Run the GC to shrink the heap as much as possible and report usage.
	Debug::log("Running GC");
	mvm_runGC(vm.get(), true);
	reportMemory();

	// Report the peak usage
	Debug::log("Peak heap used: {} bytes, peak stack used: {} bytes",
	           stats.virtualHeapHighWaterMark,
	           stats.stackHighWaterMark);
}