// WASM-PROPOSAL Phase 0 — minimal runtime loader (exit criterion).
//
// Embeds a wasm runtime through the standard wasm-c-api and links two
// modules at runtime, the way the production loader (§6) will:
//
//   1. instantiate core.wasm (owns memory/table/allocator/libc), stubbing
//      its WASI imports with named trap functions (the core stub does no
//      real I/O; any stub that actually gets called names itself)
//   2. parse unit.wasm's dylink.0 → data size/align, table slots needed
//   3. allocate __memory_base by calling core's exported malloc, and
//      __table_base by growing core's exported funcref table
//   4. build the unit's import vector: env.memory / env.__indirect_function_table /
//      env.__stack_pointer straight from core's exports; env.* functions from
//      core's exports; GOT.mem.* as mutable i32 globals holding resolved
//      addresses (self-resolved from the unit's own exports after
//      instantiation when core lacks the symbol — weak data case);
//      GOT.func.* as mutable i32 globals holding freshly grown table slots
//      pointing at core export funcrefs
//   5. instantiate unit, patch deferred GOT entries, run
//      __wasm_apply_data_relocs then __wasm_call_ctors
//   6. call uce_unit_render and read the output back out of 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)
{
	uint64_t result = 0;
	int shift = 0;
	while(true)
	{
		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");
	size_t pos = 8;
	while(pos < wasm.size())
	{
		uint8_t sec_id = wasm[pos++];
		uint64_t size = read_uleb(wasm.data(), pos);
		size_t end = pos + size;
		if(sec_id == 0)
		{
			uint64_t name_len = read_uleb(wasm.data(), pos);
			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);
					size_t sub_end = pos + sub_len;
					if(sub == 1) // WASM_DYLINK_MEM_INFO
					{
						info.mem_size = read_uleb(wasm.data(), pos);
						info.mem_align = read_uleb(wasm.data(), pos);
						info.table_size = read_uleb(wasm.data(), pos);
						info.table_align = read_uleb(wasm.data(), pos);
						info.found = true;
					}
					pos = sub_end;
				}
			}
		}
		pos = end;
	}
	return info;
}

// ---- 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];
	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-phase0/core.wasm";
	const char* unit_path = argc > 2 ? argv[2] : "/tmp/uce/wasm-phase0/unit.wasm";

	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());
		// named trap stub; leaks the name string, fine for a spike
		char* label = strdup((std::string(mod->data, mod->size) + "." + name).c_str());
		const wasm_functype_t* ft = wasm_externtype_as_functype_const(et);
		wasm_func_t* stub = wasm_func_new_with_env(g_store, ft, stub_callback, label, nullptr);
		CHECK(stub, "stub func creation failed for %s", label);
		core_import_externs[i] = wasm_func_as_extern(stub);
	}
	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");

	// ---- 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;
	int32_t raw = call_i32(core, "malloc", { (int32_t)(dl.mem_size + align) });
	CHECK(raw, "core malloc returned 0");
	int32_t memory_base = (raw + (align - 1)) & ~(int32_t)(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=%d __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(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);
			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(memory_base + v.of.i32);
		wasm_global_set(got, &nv);
		printf("self-resolved GOT.mem.%s = %d (offset %d)\n", nm.c_str(), memory_base + 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. render and read back --------------------------------------------
	int32_t counter_addr = 0;
	{
		wasm_global_t* cc = core.global("core_counter");
		CHECK(cc, "core_counter not exported");
		wasm_val_t v;
		wasm_global_get(cc, &v);
		counter_addr = v.of.i32;
	}
	byte_t* mem = wasm_memory_data(memory);
	int32_t counter_before;
	memcpy(&counter_before, mem + counter_addr, 4);

	call_i32(unit, "uce_unit_render");

	int32_t out_ptr = call_i32(core, "core_output_data");
	int32_t out_len = call_i32(core, "core_output_size");
	mem = wasm_memory_data(memory); // may have moved if memory grew
	int32_t counter_after;
	memcpy(&counter_after, mem + counter_addr, 4);

	printf("---- unit output (%d bytes) ----\n%.*s\n--------------------------------\n",
		out_len, out_len, mem + out_ptr);
	printf("core_counter (in linear memory): before=%d after=%d\n", counter_before, counter_after);

	bool ok = out_len > 0 &&
		memmem(mem + out_ptr, out_len, "unit-data-segment-ok", 20) &&
		memmem(mem + out_ptr, out_len, "counter=7", 9) &&
		memmem(mem + out_ptr, out_len, "mapsum=3", 8) &&
		memmem(mem + out_ptr, out_len, "core-string+unit", 16) &&
		memmem(mem + out_ptr, out_len, "[cb:42]", 7) &&
		memmem(mem + out_ptr, out_len, "[got-func-ok]", 13) &&
		memmem(mem + out_ptr, out_len, "[fn:42]", 7) &&
		counter_before == 7 && counter_after == 8;

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