// Copyright (c) 2023 l-m.dev. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module wasm

enum Section as u8 {
	custom_section
	type_section
	import_section
	function_section
	table_section
	memory_section
	global_section
	export_section
	start_section
	element_section
	code_section
	data_section
	data_count_section
}

enum Subsection as u8 {
	name_module
	name_function
	name_local
	// see: https://github.com/WebAssembly/extended-name-section
	name_label
	name_type
	name_table
	name_memory
	name_global
	name_elem
	name_data
}

pub enum NumType as u8 {
	i32_t = 0x7f
	i64_t = 0x7e
	f32_t = 0x7d
	f64_t = 0x7c
}

pub enum ValType as u8 {
	i32_t = 0x7f
	i64_t = 0x7e
	f32_t = 0x7d
	f64_t = 0x7c
	v128_t = 0x7b
	funcref_t = 0x70
	externref_t = 0x6f
}

pub enum RefType as u8 {
	funcref_t = 0x70
	externref_t = 0x6f
}

// Module contains the WebAssembly module.
// Use the `compile` method to compile the module into a pure byte array.
[heap]
pub struct Module {
mut:
	buf            []u8
	functypes      []FuncType
	functions      map[string]Function
	globals        []Global
	memory         ?Memory
	start          ?string
	fn_imports     []FunctionImport
	global_imports []GlobalImport
	segments       []DataSegment
	debug          bool
	mod_name       ?string
}

struct Global {
	typ    ValType
	is_mut bool
	name   string
	export bool
mut:
	init ConstExpression
}

struct GlobalImport {
	mod    string
	name   string
	typ    ValType
	is_mut bool
}

struct FunctionImport {
	mod  string
	name string
	tidx int
}

struct Memory {
	name   string
	export bool
	min    u32
	max    ?u32
}

struct DataSegment {
	idx  ?int
	data []u8
	name ?string
}

pub type LocalIndex = int
pub type GlobalIndex = int
pub type GlobalImportIndex = int
pub type DataSegmentIndex = int

pub struct FuncType {
pub:
	parameters []ValType
	results    []ValType
	name       ?string
}

fn (mut mod Module) new_functype(ft FuncType) int {
	// interns existing types
	mut idx := mod.functypes.index(ft)

	if idx == -1 {
		idx = mod.functypes.len
		mod.functypes << ft
	}

	return idx
}

// new_function creates a function struct.
pub fn (mut mod Module) new_function(name string, parameters []ValType, results []ValType) Function {
	assert name !in mod.functions.keys()

	idx := mod.functions.len
	tidx := mod.new_functype(FuncType{parameters, results, none})

	return Function{
		name: name
		tidx: tidx
		idx: idx
		mod: mod
		locals: parameters.map(FunctionLocal{}) // specifying it's ValType doesn't matter
	}
}

// new_debug_function creates a function struct with extra debug information.
// `argument_names` must be the same length as the parameters in the function type `typ`.
pub fn (mut mod Module) new_debug_function(name string, typ FuncType, argument_names []?string) Function {
	assert name !in mod.functions.keys()
	assert typ.parameters.len == argument_names.len

	idx := mod.functions.len
	tidx := mod.new_functype(typ)

	return Function{
		name: name
		tidx: tidx
		idx: idx
		mod: mod
		locals: argument_names.map(FunctionLocal{ name: it }) // specifying it's ValType doesn't matter
	}
}

// enable_debug sets whether to emit debug information for not.
pub fn (mut mod Module) enable_debug(mod_name ?string) {
	mod.debug = true
	mod.mod_name = mod_name
}

// assign_memory assigns memory to the current module.
pub fn (mut mod Module) assign_memory(name string, export bool, min u32, max ?u32) {
	mod.memory = Memory{
		name: name
		export: export
		min: min
		max: max
	}
}

// assign_start assigns the start function to the current module.
pub fn (mut mod Module) assign_start(name string) {
	mod.start = name
}

// new_function_import imports a new function into the current module.
pub fn (mut mod Module) new_function_import(modn string, name string, parameters []ValType, results []ValType) {
	assert !mod.fn_imports.any(it.mod == modn && it.name == name)

	tidx := mod.new_functype(FuncType{parameters, results, none})

	mod.fn_imports << FunctionImport{
		mod: modn
		name: name
		tidx: tidx
	}
}

// new_function_import_debug imports a new function into the current module with extra debug information.
pub fn (mut mod Module) new_function_import_debug(modn string, name string, typ FuncType) {
	assert !mod.fn_imports.any(it.mod == modn && it.name == name)

	tidx := mod.new_functype(typ)

	mod.fn_imports << FunctionImport{
		mod: modn
		name: name
		tidx: tidx
	}
}

// commit commits a function to the module, use `export` to export the function.
pub fn (mut mod Module) commit(func Function, export bool) {
	assert func.name !in mod.functions.keys()

	mod.functions[func.name] = Function{
		...func
		export: export
	}
}

// new_data_segment inserts a new data segment at the memory index `pos`.
// `name` is optional, it is used for debug info.
pub fn (mut mod Module) new_data_segment(name ?string, pos int, data []u8) DataSegmentIndex {
	len := mod.segments.len
	mod.segments << DataSegment{
		idx: pos
		data: data
		name: name
	}
	return len
}

// new_passive_data_segment inserts a new passive data segment.
// `name` is optional, it is used for debug info.
pub fn (mut mod Module) new_passive_data_segment(name ?string, data []u8) {
	mod.segments << DataSegment{
		data: data
		name: name
	}
}

// new_global creates a global and returns it's index.
// See `global_get`, `global_set`.
pub fn (mut mod Module) new_global(name string, export bool, typ ValType, is_mut bool, init ConstExpression) GlobalIndex {
	len := mod.globals.len
	mod.globals << Global{
		typ: typ
		is_mut: is_mut
		name: name
		export: export
		init: init
	}
	return len
}

// new_global_import imports a new global into the current module and returns it's index.
// See `global_get`, `global_set`.
pub fn (mut mod Module) new_global_import(modn string, name string, typ ValType, is_mut bool) GlobalImportIndex {
	assert !mod.fn_imports.any(it.mod == modn && it.name == name)

	len := mod.global_imports.len
	mod.global_imports << GlobalImport{
		mod: modn
		name: name
		typ: typ
		is_mut: is_mut
	}
	return len
}

// assign_global_init assigns a global with the constant expression `init`.
// See `new_global`.
pub fn (mut mod Module) assign_global_init(global GlobalIndex, init ConstExpression) {
	mod.globals[global].init = init
}