// WASM-PROPOSAL Phase 2 — minimal membrane host.
// Instantiates the statically linked core/page module, serves a UCEB1 request
// context via hostcall, invokes render, and reads the response from guest
// memory. Dynamic linking is intentionally Phase 3 work.

#include <wasm.h>

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

#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 wasm_store_t* g_store = nullptr;
static wasm_memory_t* g_memory = nullptr;
static std::string g_context;

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;
}

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;
	bool is_list = false;
	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(node.is_list ? 1 : 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("phase2.example.test")});
	params.children.push_back({"SCRIPT_URL", scalar("/spikes/wasm-phase2/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-phase2/page.uce")});
	root.children.push_back({"nested", nested});

	std::string out = "UCEB";
	out.push_back((char)1);
	encode_node(out, root);
	return out;
}

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())
	{
		results->data[0] = WASM_I32_VAL((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 <= mem_size, "ctx_read pointer outside memory");
	CHECK((size_t)ptr + g_context.size() <= mem_size, "ctx_read buffer outside memory");
	memcpy(mem + ptr, g_context.data(), g_context.size());
	results->data[0] = WASM_I32_VAL((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;
}

static wasm_trap_t* stub_callback(void* env, const wasm_val_vec_t* args, wasm_val_vec_t* results)
{
	(void)args; (void)results;
	char msg[256];
	snprintf(msg, sizeof(msg), "unimplemented 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, &message);
	wasm_byte_vec_delete(&message);
	return trap;
}

static wasm_func_t* make_func(wasm_store_t* store, std::vector<wasm_valkind_t> params, std::vector<wasm_valkind_t> results, wasm_func_callback_with_env_t cb, void* env = nullptr)
{
	wasm_valtype_vec_t ps;
	wasm_valtype_vec_new_uninitialized(&ps, params.size());
	for(size_t i = 0; i < params.size(); i++)
		ps.data[i] = wasm_valtype_new(params[i]);
	wasm_valtype_vec_t rs;
	wasm_valtype_vec_new_uninitialized(&rs, results.size());
	for(size_t i = 0; i < results.size(); i++)
		rs.data[i] = wasm_valtype_new(results[i]);
	wasm_functype_t* ft = wasm_functype_new(&ps, &rs);
	wasm_func_t* fn = wasm_func_new_with_env(store, ft, cb, env, nullptr);
	wasm_functype_delete(ft);
	return fn;
}

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);
		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);
			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);
	}
};

static int32_t call_i32(Instance& inst, const char* name)
{
	wasm_func_t* f = inst.func(name);
	CHECK(f, "missing export func %s", name);
	wasm_val_t result_buf[1] = { WASM_INIT_VAL };
	wasm_val_vec_t args = WASM_EMPTY_VEC;
	wasm_val_vec_t results = { wasm_func_result_arity(f), result_buf };
	wasm_val_vec_t no_results = WASM_EMPTY_VEC;
	wasm_trap_t* trap = wasm_func_call(f, &args, results.size ? &results : &no_results);
	if(trap)
	{
		wasm_message_t msg;
		wasm_trap_message(trap, &msg);
		FAIL("trap during %s: %.*s", name, (int)msg.size, msg.data);
	}
	return results.size ? result_buf[0].of.i32 : 0;
}

int main(int argc, char** argv)
{
	const char* core_path = argc > 1 ? argv[1] : "/tmp/uce/wasm-phase2/core.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");

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

	wasm_importtype_vec_t imports = WASM_EMPTY_VEC;
	wasm_module_imports(core.module, &imports);
	std::vector<wasm_extern_t*> externs(imports.size);
	std::vector<wasm_func_t*> owned_funcs;
	for(size_t i = 0; i < imports.size; i++)
	{
		const wasm_name_t* mod_n = wasm_importtype_module(imports.data[i]);
		const wasm_name_t* name_n = wasm_importtype_name(imports.data[i]);
		std::string mod(mod_n->data, mod_n->size);
		std::string name(name_n->data, name_n->size);
		while(!mod.empty() && mod.back() == '\0') mod.pop_back();
		while(!name.empty() && name.back() == '\0') name.pop_back();
		const wasm_externtype_t* et = wasm_importtype_type(imports.data[i]);
		CHECK(wasm_externtype_kind(et) == WASM_EXTERN_FUNC, "unexpected non-func import %s.%s", mod.c_str(), name.c_str());

		wasm_func_t* fn = nullptr;
		if(mod == "env" && name == "uce_host_ctx_read")
			fn = make_func(g_store, {WASM_I32, WASM_I32}, {WASM_I32}, host_ctx_read);
		else if(mod == "env" && name == "uce_host_log")
			fn = make_func(g_store, {WASM_I32, WASM_I32, WASM_I32}, {}, host_log);
		else
		{
			char* label = strdup((mod + "." + name).c_str());
			const wasm_functype_t* ft = wasm_externtype_as_functype_const(et);
			fn = wasm_func_new_with_env(g_store, ft, stub_callback, label, nullptr);
		}
		CHECK(fn, "failed to create import %s.%s", mod.c_str(), name.c_str());
		owned_funcs.push_back(fn);
		externs[i] = wasm_func_as_extern(fn);
	}

	wasm_extern_vec_t iv = { externs.size(), externs.data() };
	wasm_trap_t* trap = nullptr;
	core.instance = wasm_instance_new(g_store, core.module, &iv, &trap);
	if(trap)
	{
		wasm_message_t msg;
		wasm_trap_message(trap, &msg);
		FAIL("trap during instantiation: %.*s", (int)msg.size, msg.data);
	}
	CHECK(core.instance, "instantiation failed");
	core.index_exports();
	g_memory = wasm_extern_as_memory(core.by_name.count("memory") ? core.by_name["memory"] : nullptr);
	CHECK(g_memory, "core does not export memory");
	if(core.func("_initialize"))
		call_i32(core, "_initialize");

	int rc = call_i32(core, "uce_phase2_render");
	CHECK(rc == 0, "render returned %d", rc);
	int32_t data = call_i32(core, "uce_phase2_output_data");
	int32_t size = call_i32(core, "uce_phase2_output_size");
	uint8_t* mem = (uint8_t*)wasm_memory_data(g_memory);
	CHECK((size_t)data + (size_t)size <= wasm_memory_data_size(g_memory), "output outside memory");
	std::string output((const char*)mem + data, size);
	printf("---- phase2 output (%d bytes) ----\n%s", size, output.c_str());
	printf("----------------------------------\n");
	CHECK(output.find("PHASE2 PAGE OK") != std::string::npos, "missing page marker");
	CHECK(output.find("host=phase2.example.test") != std::string::npos, "missing context param");
	CHECK(output.find("answer=42") != std::string::npos, "missing nested context value");
	printf("PHASE2 EXIT CRITERION: PASS\n");
	return 0;
}