// WASM-PROPOSAL Phase 3 — dynamic loader + membrane scaffold.
//
// Embeds a wasm runtime through the standard wasm-c-api and links a core
// workspace module with a generated-shape PIC UCE side module:
//
//   1. instantiate core.wasm (owns memory/table/allocator/Request/DValue)
//   2. parse page.wasm's dylink.0 → data size/align, table slots needed
//   3. allocate __memory_base by calling core malloc, and __table_base by
//      appending to the shared funcref table
//   4. build the unit's import vector: env.memory / env.__indirect_function_table /
//      env.__stack_pointer from core; env.* functions from core exports;
//      GOT.mem.* and GOT.func.* through the same rules as Phase 0
//   5. instantiate unit, patch deferred GOT entries, run relocations/ctors
//   6. prepare a UCEB1 request context in core, set the unit's Request*, call
//      __uce_render, then read output back out of core-owned linear memory
//
// Everything here is deliberately validation-grade: fail loudly, never guess.

#include <wasm.h>

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include <map>

#define FAIL(...) do { fprintf(stderr, "FAIL: " __VA_ARGS__); fprintf(stderr, "\n"); exit(1); } while(0)
#define CHECK(cond, ...) do { if(!(cond)) FAIL(__VA_ARGS__); } while(0)

static std::vector<uint8_t> read_file(const char* fn)
{
	FILE* f = fopen(fn, "rb");
	CHECK(f, "cannot open %s", fn);
	fseek(f, 0, SEEK_END);
	long n = ftell(f);
	fseek(f, 0, SEEK_SET);
	std::vector<uint8_t> buf(n);
	CHECK(fread(buf.data(), 1, n, f) == (size_t)n, "short read on %s", fn);
	fclose(f);
	return buf;
}

// ---- minimal dylink.0 parser -------------------------------------------

struct DylinkInfo
{
	uint32_t mem_size = 0;
	uint32_t mem_align = 0;   // power of 2
	uint32_t table_size = 0;
	uint32_t table_align = 0;
	bool found = false;
};

static uint64_t read_uleb(const uint8_t* buf, size_t& pos, size_t end)
{
	uint64_t result = 0;
	int shift = 0;
	while(true)
	{
		CHECK(pos < end, "truncated uleb in wasm custom section");
		CHECK(shift < 64, "oversized uleb in wasm custom section");
		uint8_t b = buf[pos++];
		result |= (uint64_t)(b & 0x7f) << shift;
		if(!(b & 0x80))
			return result;
		shift += 7;
	}
}

static DylinkInfo parse_dylink(const std::vector<uint8_t>& wasm)
{
	DylinkInfo info;
	CHECK(wasm.size() >= 8 && !memcmp(wasm.data(), "\0asm", 4), "not a wasm module");
	CHECK(wasm[4] == 1 && wasm[5] == 0 && wasm[6] == 0 && wasm[7] == 0, "unsupported wasm binary version");
	size_t pos = 8;
	while(pos < wasm.size())
	{
		uint8_t sec_id = wasm[pos++];
		uint64_t size = read_uleb(wasm.data(), pos, wasm.size());
		CHECK(size <= wasm.size() - pos, "wasm section exceeds file size");
		size_t end = pos + size;
		if(sec_id == 0)
		{
			uint64_t name_len = read_uleb(wasm.data(), pos, end);
			CHECK(name_len <= end - pos, "custom section name exceeds section size");
			std::string name((const char*)wasm.data() + pos, name_len);
			pos += name_len;
			if(name == "dylink.0")
			{
				while(pos < end)
				{
					uint8_t sub = wasm[pos++];
					uint64_t sub_len = read_uleb(wasm.data(), pos, end);
					CHECK(sub_len <= end - pos, "dylink subsection exceeds section size");
					size_t sub_end = pos + sub_len;
					if(sub == 1) // WASM_DYLINK_MEM_INFO
					{
						info.mem_size = read_uleb(wasm.data(), pos, sub_end);
						info.mem_align = read_uleb(wasm.data(), pos, sub_end);
						info.table_size = read_uleb(wasm.data(), pos, sub_end);
						info.table_align = read_uleb(wasm.data(), pos, sub_end);
						CHECK(!info.found, "duplicate dylink.0 mem_info subsection");
						CHECK(info.mem_align < 31 && info.table_align < 31, "unsupported dylink alignment");
						info.found = true;
					}
					pos = sub_end;
				}
			}
		}
		pos = end;
	}
	return info;
}

// ---- tiny UCEB1 request-context encoder ---------------------------------

static void append_varuint(std::string& out, uint64_t value)
{
	while(value >= 0x80)
	{
		out.push_back((char)((value & 0x7f) | 0x80));
		value >>= 7;
	}
	out.push_back((char)value);
}

struct Node
{
	std::string scalar;
	std::vector<std::pair<std::string, Node>> children;
};

static Node scalar(const char* value)
{
	Node n;
	n.scalar = value;
	return n;
}

static void encode_node(std::string& out, const Node& node)
{
	out.push_back(0);
	append_varuint(out, node.scalar.size());
	out.append(node.scalar);
	append_varuint(out, node.children.size());
	for(const auto& child : node.children)
	{
		append_varuint(out, child.first.size());
		out.append(child.first);
		encode_node(out, child.second);
	}
}

static std::string make_context()
{
	Node params;
	params.children.push_back({"HTTP_HOST", scalar("phase3.example.test")});
	params.children.push_back({"SCRIPT_URL", scalar("/spikes/wasm-phase3/page.uce")});
	Node nested;
	nested.children.push_back({"answer", scalar("42")});
	Node root;
	root.children.push_back({"params", params});
	root.children.push_back({"route", scalar("/spikes/wasm-phase3/page.uce")});
	root.children.push_back({"nested", nested});
	std::string out = "UCEB";
	out.push_back((char)1);
	encode_node(out, root);
	return out;
}

// ---- host imports --------------------------------------------------------

static wasm_memory_t* g_memory = nullptr;
static std::string g_context;

static void set_i32_result(wasm_val_t* result, int32_t value)
{
	result->kind = WASM_I32;
	result->of.i32 = value;
}

static wasm_trap_t* host_ctx_read(void* env, const wasm_val_vec_t* args, wasm_val_vec_t* results)
{
	(void)env;
	uint32_t ptr = args->data[0].of.i32;
	uint32_t cap = args->data[1].of.i32;
	if(ptr == 0 || cap == 0 || cap < g_context.size())
	{
		set_i32_result(&results->data[0], (int32_t)g_context.size());
		return nullptr;
	}
	CHECK(g_memory, "host ctx_read called before memory export was captured");
	uint8_t* mem = (uint8_t*)wasm_memory_data(g_memory);
	size_t mem_size = wasm_memory_data_size(g_memory);
	CHECK((size_t)ptr + g_context.size() <= mem_size, "ctx_read buffer outside memory");
	memcpy(mem + ptr, g_context.data(), g_context.size());
	set_i32_result(&results->data[0], (int32_t)g_context.size());
	return nullptr;
}

static wasm_trap_t* host_log(void* env, const wasm_val_vec_t* args, wasm_val_vec_t* results)
{
	(void)env; (void)results;
	int level = args->data[0].of.i32;
	uint32_t ptr = args->data[1].of.i32;
	uint32_t len = args->data[2].of.i32;
	uint8_t* mem = (uint8_t*)wasm_memory_data(g_memory);
	size_t mem_size = wasm_memory_data_size(g_memory);
	if((size_t)ptr + len <= mem_size)
		fprintf(stderr, "[guest log %d] %.*s\n", level, (int)len, (const char*)mem + ptr);
	return nullptr;
}

// ---- named trap stubs for unsatisfied (WASI) imports --------------------

static wasm_store_t* g_store_for_stubs = nullptr;

static wasm_trap_t* stub_callback(void* env, const wasm_val_vec_t* args, wasm_val_vec_t* results)
{
	(void)args; (void)results;
	fprintf(stderr, "[stub called: %s]\n", (const char*)env);
	char msg[256];
	snprintf(msg, sizeof(msg), "unimplemented host import called: %s", (const char*)env);
	wasm_message_t message;
	wasm_byte_vec_new(&message, strlen(msg) + 1, msg);
	wasm_trap_t* trap = wasm_trap_new(g_store_for_stubs, &message);
	wasm_byte_vec_delete(&message);
	return trap;
}

// ---- instance wrapper: name → extern map --------------------------------

struct Instance
{
	wasm_module_t* module = nullptr;
	wasm_instance_t* instance = nullptr;
	wasm_extern_vec_t exports = WASM_EMPTY_VEC;
	std::map<std::string, wasm_extern_t*> by_name;

	void index_exports()
	{
		wasm_exporttype_vec_t types = WASM_EMPTY_VEC;
		wasm_module_exports(module, &types);
		wasm_instance_exports(instance, &exports);
		CHECK(types.size == exports.size, "export type/extern count mismatch");
		for(size_t i = 0; i < types.size; i++)
		{
			const wasm_name_t* nm = wasm_exporttype_name(types.data[i]);
			std::string key(nm->data, nm->size);
			// some wasm-c-api impls include the trailing NUL in name size
			while(!key.empty() && key.back() == '\0')
				key.pop_back();
			by_name[key] = exports.data[i];
		}
		wasm_exporttype_vec_delete(&types);
	}

	wasm_func_t* func(const char* name)
	{
		auto it = by_name.find(name);
		return it == by_name.end() ? nullptr : wasm_extern_as_func(it->second);
	}
	wasm_global_t* global(const char* name)
	{
		auto it = by_name.find(name);
		return it == by_name.end() ? nullptr : wasm_extern_as_global(it->second);
	}
};

static wasm_store_t* g_store = nullptr;

static void report_trap(wasm_trap_t* trap, const char* what)
{
	if(!trap)
		return;
	wasm_message_t msg;
	wasm_trap_message(trap, &msg);
	FAIL("trap during %s: %.*s", what, (int)msg.size, msg.data);
}

static int32_t call_i32(Instance& inst, const char* name, std::vector<int32_t> argv = {})
{
	wasm_func_t* f = inst.func(name);
	CHECK(f, "missing export func %s", name);
	wasm_val_t args_buf[4];
	CHECK(argv.size() <= 4, "call_i32 argv overflow for %s", name);
	for(size_t i = 0; i < argv.size(); i++)
		args_buf[i] = WASM_I32_VAL(argv[i]);
	wasm_val_t results_buf[1] = { WASM_INIT_VAL };
	wasm_val_vec_t args = { argv.size(), args_buf };
	wasm_val_vec_t results = { 1, results_buf };
	size_t result_arity = wasm_func_result_arity(f);
	wasm_val_vec_t no_results = WASM_EMPTY_VEC;
	wasm_trap_t* trap = wasm_func_call(f, &args, result_arity ? &results : &no_results);
	report_trap(trap, name);
	return result_arity ? results_buf[0].of.i32 : 0;
}

static wasm_global_t* make_i32_global(int32_t value, wasm_mutability_t mut)
{
	wasm_globaltype_t* gt = wasm_globaltype_new(wasm_valtype_new(WASM_I32), mut);
	wasm_val_t val = WASM_I32_VAL(value);
	wasm_global_t* g = wasm_global_new(g_store, gt, &val);
	CHECK(g, "wasm_global_new failed (host-created globals unsupported?)");
	wasm_globaltype_delete(gt);
	return g;
}

int main(int argc, char** argv)
{
	const char* core_path = argc > 1 ? argv[1] : "/tmp/uce/wasm-phase3/core.wasm";
	const char* unit_path = argc > 2 ? argv[2] : "/tmp/uce/wasm-phase3/page.wasm";
	g_context = make_context();

	wasm_engine_t* engine = wasm_engine_new();
	CHECK(engine, "engine");
	g_store = wasm_store_new(engine);
	CHECK(g_store, "store");
	g_store_for_stubs = g_store;

	// ---- 1. core module ---------------------------------------------------
	std::vector<uint8_t> core_bytes = read_file(core_path);
	wasm_byte_vec_t core_bv;
	wasm_byte_vec_new(&core_bv, core_bytes.size(), (const char*)core_bytes.data());
	Instance core;
	core.module = wasm_module_new(g_store, &core_bv);
	wasm_byte_vec_delete(&core_bv);
	CHECK(core.module, "core module load failed");

	// the shared funcref table is host-created (core links with --import-table)
	// because WAMR cannot grow a table from the host: size it up front as
	// core's declared minimum (= its own elem needs) plus headroom for unit
	// module table regions. __table_base allocation bumps from the minimum.
	wasm_table_t* table = nullptr;
	uint32_t table_next_free = 0;
	const uint32_t TABLE_HEADROOM = 2048;

	wasm_importtype_vec_t core_imports = WASM_EMPTY_VEC;
	wasm_module_imports(core.module, &core_imports);
	std::vector<wasm_extern_t*> core_import_externs(core_imports.size);
	for(size_t i = 0; i < core_imports.size; i++)
	{
		const wasm_name_t* mod = wasm_importtype_module(core_imports.data[i]);
		const wasm_name_t* nm = wasm_importtype_name(core_imports.data[i]);
		const wasm_externtype_t* et = wasm_importtype_type(core_imports.data[i]);
		std::string name(nm->data, nm->size);
		if(wasm_externtype_kind(et) == WASM_EXTERN_TABLE)
		{
			CHECK(name.rfind("__indirect_function_table", 0) == 0,
				"unexpected core table import %s", name.c_str());
			const wasm_tabletype_t* tt = wasm_externtype_as_tabletype_const(et);
			uint32_t core_min = wasm_tabletype_limits(tt)->min;
			wasm_limits_t limits = { core_min + TABLE_HEADROOM, core_min + TABLE_HEADROOM };
			wasm_tabletype_t* host_tt = wasm_tabletype_new(wasm_valtype_new(WASM_FUNCREF), &limits);
			table = wasm_table_new(g_store, host_tt, nullptr);
			CHECK(table, "wasm_table_new failed (host-created tables unsupported?)");
			wasm_tabletype_delete(host_tt);
			table_next_free = core_min;
			printf("host table created: size=%u, core region=[0,%u)\n", core_min + TABLE_HEADROOM, core_min);
			core_import_externs[i] = wasm_table_as_extern(table);
			continue;
		}
		CHECK(wasm_externtype_kind(et) == WASM_EXTERN_FUNC,
			"core has unexpected non-func import %.*s.%s",
			(int)mod->size, mod->data, name.c_str());
		const wasm_functype_t* ft = wasm_externtype_as_functype_const(et);
		wasm_func_t* import_func = nullptr;
		std::string mod_name(mod->data, mod->size);
		while(!mod_name.empty() && mod_name.back() == '\0') mod_name.pop_back();
		while(!name.empty() && name.back() == '\0') name.pop_back();
		if(mod_name == "env" && name == "uce_host_ctx_read")
			import_func = wasm_func_new_with_env(g_store, ft, host_ctx_read, nullptr, nullptr);
		else if(mod_name == "env" && name == "uce_host_log")
			import_func = wasm_func_new_with_env(g_store, ft, host_log, nullptr, nullptr);
		else
		{
			// named trap stub; leaks the name string, fine for a spike
			char* label = strdup((mod_name + "." + name).c_str());
			import_func = wasm_func_new_with_env(g_store, ft, stub_callback, label, nullptr);
		}
		CHECK(import_func, "stub func creation failed for %s", name.c_str());
		core_import_externs[i] = wasm_func_as_extern(import_func);
	}
	CHECK(table, "core does not import __indirect_function_table — rebuild with --import-table");
	wasm_extern_vec_t core_iv = { core_import_externs.size(), core_import_externs.data() };
	wasm_trap_t* trap = nullptr;
	core.instance = wasm_instance_new(g_store, core.module, &core_iv, &trap);
	report_trap(trap, "core instantiation");
	CHECK(core.instance, "core instantiation failed");
	core.index_exports();
	printf("core instantiated: %zu exports\n", core.by_name.size());

	// reactor init (runs ctors). Wasmtime does not auto-run _initialize;
	// WAMR does — and calling it twice trips wasi-libc's double-init guard
	// (__builtin_trap → "unreachable"), which cost us a debugging round.
	if(core.func("_initialize"))
		call_i32(core, "_initialize");

	wasm_memory_t* memory = wasm_extern_as_memory(core.by_name.count("memory") ? core.by_name["memory"] : nullptr);
	CHECK(memory, "core does not export memory");
	g_memory = memory;

	// ---- 2./3. dylink + base allocation ------------------------------------
	std::vector<uint8_t> unit_bytes = read_file(unit_path);
	DylinkInfo dl = parse_dylink(unit_bytes);
	CHECK(dl.found, "unit has no dylink.0 mem_info");
	printf("dylink.0: memsize=%u memalign=2^%u tablesize=%u\n", dl.mem_size, dl.mem_align, dl.table_size);

	uint32_t align = 1u << dl.mem_align;
	uint32_t raw = (uint32_t)call_i32(core, "malloc", { (int32_t)(dl.mem_size + align) });
	CHECK(raw, "core malloc returned 0");
	uint32_t memory_base = (raw + (align - 1)) & ~(align - 1);

	uint32_t table_base = table_next_free;
	table_next_free += dl.table_size;
	CHECK(table_next_free <= wasm_table_size(table), "table headroom exhausted");
	printf("bases: __memory_base=%u __table_base=%u (table size %u)\n",
		memory_base, table_base, wasm_table_size(table));

	// ---- 4. unit import resolution -----------------------------------------
	wasm_byte_vec_t unit_bv;
	wasm_byte_vec_new(&unit_bv, unit_bytes.size(), (const char*)unit_bytes.data());
	Instance unit;
	unit.module = wasm_module_new(g_store, &unit_bv);
	wasm_byte_vec_delete(&unit_bv);
	CHECK(unit.module, "unit module load failed");

	wasm_importtype_vec_t unit_imports = WASM_EMPTY_VEC;
	wasm_module_imports(unit.module, &unit_imports);
	std::vector<wasm_extern_t*> unit_import_externs(unit_imports.size);
	// GOT.mem entries that must be self-resolved from the unit's own exports
	// after instantiation (weak data the core does not define)
	std::vector<std::pair<std::string, wasm_global_t*>> deferred_got;

	for(size_t i = 0; i < unit_imports.size; i++)
	{
		const wasm_name_t* mod_n = wasm_importtype_module(unit_imports.data[i]);
		const wasm_name_t* nm_n = wasm_importtype_name(unit_imports.data[i]);
		std::string mod(mod_n->data, mod_n->size);
		std::string nm(nm_n->data, nm_n->size);
		const wasm_externtype_t* et = wasm_importtype_type(unit_imports.data[i]);
		wasm_externkind_t kind = wasm_externtype_kind(et);
		wasm_extern_t* resolved = nullptr;

		if(mod == "env" && nm == "memory")
			resolved = wasm_memory_as_extern(memory);
		else if(mod == "env" && nm == "__indirect_function_table")
			resolved = wasm_table_as_extern(table);
		else if(mod == "env" && nm == "__stack_pointer")
		{
			CHECK(core.global("__stack_pointer"), "core does not export __stack_pointer");
			resolved = core.by_name["__stack_pointer"];
		}
		else if(mod == "env" && nm == "__memory_base")
			resolved = wasm_global_as_extern(make_i32_global((int32_t)memory_base, WASM_CONST));
		else if(mod == "env" && nm == "__table_base")
			resolved = wasm_global_as_extern(make_i32_global((int32_t)table_base, WASM_CONST));
		else if(mod == "env" && kind == WASM_EXTERN_FUNC)
		{
			auto it = core.by_name.find(nm);
			CHECK(it != core.by_name.end(), "unresolved unit func import env.%s", nm.c_str());
			resolved = it->second;
		}
		else if(mod == "GOT.mem")
		{
			wasm_global_t* src = core.global(nm.c_str());
			if(src)
			{
				wasm_val_t v;
				wasm_global_get(src, &v);
				resolved = wasm_global_as_extern(make_i32_global(v.of.i32, WASM_VAR));
			}
			else
			{
				// provisional 0; patched from the unit's own export post-instantiation
				wasm_global_t* g = make_i32_global(0, WASM_VAR);
				deferred_got.push_back({ nm, g });
				resolved = wasm_global_as_extern(g);
			}
		}
		else if(mod == "GOT.func")
		{
			// resolved guest-side: the core exports a helper returning
			// (intptr_t)&func, which on wasm is the function's table index —
			// no host-side funcref injection needed (WAMR forbids it anyway)
			std::string helper = "core_table_index_of_" + nm;
			CHECK(core.func(helper.c_str()), "no GOT.func resolver %s in core", helper.c_str());
			int32_t slot = call_i32(core, helper.c_str());
			CHECK(slot > 0, "GOT.func resolver %s returned %d", helper.c_str(), slot);
			printf("resolved GOT.func.%s = table[%d]\n", nm.c_str(), slot);
			resolved = wasm_global_as_extern(make_i32_global(slot, WASM_VAR));
		}
		CHECK(resolved, "unhandled unit import %s.%s (kind %d)", mod.c_str(), nm.c_str(), (int)kind);
		unit_import_externs[i] = resolved;
	}

	// ---- 5. instantiate unit, patch GOT, run init ---------------------------
	wasm_extern_vec_t unit_iv = { unit_import_externs.size(), unit_import_externs.data() };
	trap = nullptr;
	unit.instance = wasm_instance_new(g_store, unit.module, &unit_iv, &trap);
	report_trap(trap, "unit instantiation");
	CHECK(unit.instance, "unit instantiation failed");
	unit.index_exports();
	printf("unit instantiated: %zu exports\n", unit.by_name.size());

	for(auto& [nm, got] : deferred_got)
	{
		wasm_global_t* own = unit.global(nm.c_str());
		CHECK(own, "GOT.mem.%s defined neither by core nor by unit", nm.c_str());
		wasm_val_t v;
		wasm_global_get(own, &v);
		// dylink ABI: a PIC module's exported data symbols are offsets relative
		// to its __memory_base; the linker must add the base when resolving
		wasm_val_t nv = WASM_I32_VAL((int32_t)(memory_base + (uint32_t)v.of.i32));
		wasm_global_set(got, &nv);
		printf("self-resolved GOT.mem.%s = %u (offset %d)\n", nm.c_str(), memory_base + (uint32_t)v.of.i32, v.of.i32);
	}

	if(unit.func("__wasm_apply_data_relocs"))
		call_i32(unit, "__wasm_apply_data_relocs");
	if(unit.func("__wasm_call_ctors"))
		call_i32(unit, "__wasm_call_ctors");

	// ---- 6. membrane prepare, render and read back ---------------------------
	int32_t rc = call_i32(core, "uce_phase3_prepare");
	CHECK(rc == 0, "uce_phase3_prepare returned %d", rc);
	int32_t request_ptr = call_i32(core, "uce_phase3_request");
	CHECK(request_ptr != 0, "core returned null Request*");
	call_i32(unit, "__uce_set_current_request", { request_ptr });
	call_i32(unit, "__uce_render", { request_ptr });
	call_i32(core, "uce_phase3_finish");

	int32_t out_ptr = call_i32(core, "uce_phase3_output_data");
	int32_t out_len = call_i32(core, "uce_phase3_output_size");
	byte_t* mem = wasm_memory_data(memory); // may have moved if memory grew

	printf("---- phase3 unit output (%d bytes) ----\n%.*s\n--------------------------------------\n",
		out_len, out_len, mem + out_ptr);

	bool ok = out_len > 0 &&
		memmem(mem + out_ptr, out_len, "PHASE3 PAGE OK", 14) &&
		memmem(mem + out_ptr, out_len, "host=phase3.example.test", 24) &&
		memmem(mem + out_ptr, out_len, "answer=42", 9) &&
		memmem(mem + out_ptr, out_len, "each=answer:42", 14) &&
		memmem(mem + out_ptr, out_len, "self-got=42", 11) &&
		memmem(mem + out_ptr, out_len, "map=wise", 8) &&
		memmem(mem + out_ptr, out_len, "callback=42", 11) &&
		memmem(mem + out_ptr, out_len, "got-func-ok", 11) &&
		memmem(mem + out_ptr, out_len, "fn=42", 5);

	printf(ok ? "PHASE3 EXIT CRITERION: PASS\n" : "PHASE3 EXIT CRITERION: FAIL (see output above)\n");
	return ok ? 0 : 1;
}