module wasm

// encode.v is the self-contained WebAssembly binary encoder used by the v3
// wasm backend. It mirrors the role of the arm64 backend's asm.v + macho.v:
// gen.v walks the flat AST and drives this module, which knows nothing about V
// and only emits raw WASM bytecode (LEB128 + sections). No external module is
// imported, so v3 stays self-contained and self-hostable.

// WASM value types.
pub const valtype_i32 = u8(0x7f)
pub const valtype_i64 = u8(0x7e)
pub const valtype_f32 = u8(0x7d)
pub const valtype_f64 = u8(0x7c)

// Export kinds.
pub const export_func = u8(0x00)
pub const export_mem = u8(0x02)

// FuncType is a single `(params) -> (results)` signature.
struct FuncType {
    params  []u8
    results []u8
}

// ImportFunc is an imported function (only the WASI import is used for now).
struct ImportFunc {
    module   string
    name     string
    type_idx int
}

// Func is a locally defined function. `code` already ends with the `end` (0x0b)
// opcode; `locals` lists the value type of each local declared beyond params.
struct Func {
    type_idx int
    locals   []u8
    code     []u8
}

struct Export {
    name  string
    kind  u8
    index int
}

// DataSeg is an active data segment placed at a fixed linear-memory offset.
struct DataSeg {
    offset int
    bytes  []u8
}

// Global is a module-level mutable global variable.
struct Global {
    valtype u8
    init    []u8 // constant init expression, ending with `end` (0x0b)
}

@[heap]
pub struct Module {
mut:
    types    []FuncType
    imports  []ImportFunc
    funcs    []Func
    exports  []Export
    datas    []DataSeg
    globals  []Global
    mem_min  int = 2
    n_import int
}

pub fn Module.new() &Module {
    return &Module{}
}

// add_type registers a signature, deduplicating identical ones.
pub fn (mut m Module) add_type(params []u8, results []u8) int {
    for i, t in m.types {
        if t.params == params && t.results == results {
            return i
        }
    }
    m.types << FuncType{
        params:  params.clone()
        results: results.clone()
    }
    return m.types.len - 1
}

// add_import_func registers an imported function and returns its function index.
// Imports must be added before any defined function so indices stay stable.
pub fn (mut m Module) add_import_func(module_ string, name string, type_idx int) int {
    m.imports << ImportFunc{
        module:   module_
        name:     name
        type_idx: type_idx
    }
    m.n_import = m.imports.len
    return m.imports.len - 1
}

// add_global appends a mutable global and returns its index.
pub fn (mut m Module) add_global(valtype u8, init []u8) int {
    m.globals << Global{
        valtype: valtype
        init:    init.clone()
    }
    return m.globals.len - 1
}

// reserve_func_index returns the function index a later add_func will occupy.
pub fn (m &Module) reserve_func_index(defined_so_far int) int {
    return m.n_import + defined_so_far
}

// add_func appends a defined function and returns its global function index.
pub fn (mut m Module) add_func(type_idx int, locals []u8, code []u8) int {
    m.funcs << Func{
        type_idx: type_idx
        locals:   locals.clone()
        code:     code.clone()
    }
    return m.n_import + m.funcs.len - 1
}

pub fn (mut m Module) add_export(name string, kind u8, index int) {
    m.exports << Export{
        name:  name
        kind:  kind
        index: index
    }
}

pub fn (mut m Module) add_data(offset int, bytes []u8) {
    m.datas << DataSeg{
        offset: offset
        bytes:  bytes.clone()
    }
}

pub fn (mut m Module) set_mem_min(pages int) {
    m.mem_min = pages
}

// ---- LEB128 ----

fn leb_u(mut out []u8, val_ u64) {
    mut val := val_
    for {
        mut b := u8(val & 0x7f)
        val >>= 7
        if val != 0 {
            b |= 0x80
        }
        out << b
        if val == 0 {
            break
        }
    }
}

fn leb_i(mut out []u8, val_ i64) {
    mut val := val_
    for {
        mut b := u8(val & 0x7f)
        val >>= 7
        done := (val == 0 && (b & 0x40) == 0) || (val == -1 && (b & 0x40) != 0)
        if !done {
            b |= 0x80
        }
        out << b
        if done {
            break
        }
    }
}

fn write_name(mut out []u8, name string) {
    leb_u(mut out, u64(name.len))
    for c in name {
        out << c
    }
}

// section wraps a section body with its id byte and LEB-encoded length.
fn section(mut out []u8, id u8, body []u8) {
    if body.len == 0 {
        return
    }
    out << id
    leb_u(mut out, u64(body.len))
    out << body
}

// compress_locals run-length-encodes the per-local value types into the wasm
// locals vector: count-of-groups, then (count, valtype) groups.
fn compress_locals(locals []u8) []u8 {
    mut groups := [][]int{}
    for vt in locals {
        if groups.len > 0 && groups[groups.len - 1][1] == int(vt) {
            groups[groups.len - 1][0]++
        } else {
            groups << [1, int(vt)]
        }
    }
    mut out := []u8{}
    leb_u(mut out, u64(groups.len))
    for g in groups {
        leb_u(mut out, u64(g[0]))
        out << u8(g[1])
    }
    return out
}

// compile serializes the whole module into a `.wasm` byte buffer.
pub fn (m &Module) compile() []u8 {
    mut out := []u8{}
    // magic + version
    out << [u8(0x00), 0x61, 0x73, 0x6d]
    out << [u8(0x01), 0x00, 0x00, 0x00]

    // type section (1)
    mut tsec := []u8{}
    leb_u(mut tsec, u64(m.types.len))
    for t in m.types {
        tsec << 0x60
        leb_u(mut tsec, u64(t.params.len))
        tsec << t.params
        leb_u(mut tsec, u64(t.results.len))
        tsec << t.results
    }
    section(mut out, 0x01, tsec)

    // import section (2)
    if m.imports.len > 0 {
        mut isec := []u8{}
        leb_u(mut isec, u64(m.imports.len))
        for imp in m.imports {
            write_name(mut isec, imp.module)
            write_name(mut isec, imp.name)
            isec << 0x00 // import kind: func
            leb_u(mut isec, u64(imp.type_idx))
        }
        section(mut out, 0x02, isec)
    }

    // function section (3)
    mut fsec := []u8{}
    leb_u(mut fsec, u64(m.funcs.len))
    for f in m.funcs {
        leb_u(mut fsec, u64(f.type_idx))
    }
    section(mut out, 0x03, fsec)

    // memory section (5)
    mut msec := []u8{}
    leb_u(mut msec, 1) // one memory
    msec << 0x00 // limits: min only
    leb_u(mut msec, u64(m.mem_min))
    section(mut out, 0x05, msec)

    // global section (6)
    if m.globals.len > 0 {
        mut gsec := []u8{}
        leb_u(mut gsec, u64(m.globals.len))
        for gl in m.globals {
            gsec << gl.valtype
            gsec << 0x01 // mutable
            gsec << gl.init
        }
        section(mut out, 0x06, gsec)
    }

    // export section (7)
    mut esec := []u8{}
    leb_u(mut esec, u64(m.exports.len))
    for e in m.exports {
        write_name(mut esec, e.name)
        esec << e.kind
        leb_u(mut esec, u64(e.index))
    }
    section(mut out, 0x07, esec)

    // code section (10)
    mut csec := []u8{}
    leb_u(mut csec, u64(m.funcs.len))
    for f in m.funcs {
        mut body := compress_locals(f.locals)
        body << f.code
        leb_u(mut csec, u64(body.len))
        csec << body
    }
    section(mut out, 0x0a, csec)

    // data section (11)
    if m.datas.len > 0 {
        mut dsec := []u8{}
        leb_u(mut dsec, u64(m.datas.len))
        for d in m.datas {
            dsec << 0x00 // active segment, memory 0
            dsec << 0x41 // i32.const offset
            leb_i(mut dsec, i64(d.offset))
            dsec << 0x0b // end of offset expr
            leb_u(mut dsec, u64(d.bytes.len))
            dsec << d.bytes
        }
        section(mut out, 0x0b, dsec)
    }

    return out
}