// 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

import wasm
import v.ast
import v.gen.wasm.serialise
import encoding.binary

pub struct Var {
    name string
mut:
    typ        ast.Type
    idx        wasm.LocalIndex
    is_address bool
    is_global  bool
    g_idx      wasm.GlobalIndex
    offset     int
}

pub fn (mut g Gen) get_var_from_ident(ident ast.Ident) Var {
    mut obj := ident.obj
    if obj !in [ast.Var, ast.ConstField, ast.GlobalField, ast.AsmRegister] {
        obj = ident.scope.find(ident.name) or { g.w_error('unknown variable ${ident.name}') }
    }

    match mut obj {
        ast.Var {
            if g.local_vars.len == 0 {
                g.w_error('get_var_from_ident: g.local_vars.len == 0')
            }
            mut c := g.local_vars.len
            for {
                c--
                if g.local_vars[c].name == obj.name {
                    return g.local_vars[c]
                }
                if c == 0 {
                    break
                }
            }
            g.w_error('get_var_from_ident: unreachable, variable not found')
        }
        ast.ConstField {
            if gbl := g.global_vars[obj.name] {
                return gbl.v
            }
            gbl := g.new_global(obj.name, obj.typ, obj.expr, false)
            g.global_vars[obj.name] = gbl
            return gbl.v
        }
        ast.GlobalField {
            if gbl := g.global_vars[obj.name] {
                return gbl.v
            }
            gbl := g.new_global(obj.name, obj.typ, obj.expr, true)
            g.global_vars[obj.name] = gbl
            return gbl.v
        }
        else {
            g.w_error('unsupported variable type type:${obj} name:${ident.name}')
        }
    }
}

pub fn (mut g Gen) get_var_from_expr(node ast.Expr) ?Var {
    match node {
        ast.Ident {
            return g.get_var_from_ident(node)
        }
        ast.ParExpr {
            return g.get_var_from_expr(node.expr)
        }
        ast.CastExpr {
            // For cast expressions like &u32(), we need to look through
            // the cast to get the underlying variable because WASM don't have
            // really pointers like in C
            return g.get_var_from_expr(node.expr)
        }
        ast.SelectorExpr {
            mut addr := g.get_var_from_expr(node.expr) or {
                // if place {
                //     g.field_offset(node.expr_type, node.field_name)
                // }
                return none
            }
            addr.is_address = true

            offset := g.get_field_offset(node.expr_type, node.field_name)
            return g.offset(addr, node.typ, offset)
        }
        else {
            // g.w_error('get_var_from_expr: unexpected `${node.type_name()}`')
            return none
        }
    }
}

// ONLY call this with the LHS of an assign, an lvalue
// use get_var_from_expr in all other cases
pub fn (mut g Gen) get_var_or_make_from_expr(node ast.Expr, typ ast.Type) Var {
    if v := g.get_var_from_expr(node) {
        return v
    }

    mut v := g.new_local('__tmp', ast.voidptr_type)
    g.needs_address = true
    {
        g.set_with_expr(node, v)
    }
    g.needs_address = false

    v.typ = typ
    v.is_address = true

    return v
}

pub fn (mut g Gen) bp() wasm.LocalIndex {
    if g.bp_idx == -1 {
        g.bp_idx = g.func.new_local_named(.i32_t, '__vbp')
    }
    return g.bp_idx
}

pub fn (mut g Gen) sp() wasm.GlobalIndex {
    if sp := g.sp_global {
        return sp
    }
    // (64KiB - 1KiB) of stack space, grows downwards.
    // Memory addresses from 0 to 1024 are forbidden.
    g.sp_global = g.mod.new_global('__vsp', false, .i32_t, true, wasm.constexpr_value(0))
    return g.sp()
}

pub fn (mut g Gen) hp() wasm.GlobalIndex {
    if hp := g.heap_base {
        return hp
    }
    hp := g.mod.new_global('__heap_base', false, .i32_t, false, wasm.constexpr_value(0))
    g.heap_base = hp
    return hp
}

pub fn (mut g Gen) new_local(name string, typ_ ast.Type) Var {
    mut typ := typ_
    ts := g.table.sym(typ)

    match ts.info {
        ast.Enum {
            typ = ts.info.typ
        }
        ast.Alias {
            typ = ts.info.parent_type
        }
        else {}
    }

    is_address := !g.is_pure_type(typ)
    wtyp := g.get_wasm_type(typ)

    mut v := Var{
        name:       name
        typ:        typ
        is_address: is_address
    }

    if !is_address {
        v.idx = g.func.new_local_named(wtyp, g.dbg_type_name(name, typ_))
        g.local_vars << v
        return v
    }
    v.idx = g.bp()

    // allocate memory, then assign an offset
    //
    match ts.info {
        ast.Struct, ast.ArrayFixed {
            size, align := g.pool.type_size(typ)
            padding := calc_padding(g.stack_frame, align)
            address := g.stack_frame
            g.stack_frame += size + padding

            v.offset = address
        }
        else {
            g.w_error('new_local: type `${*ts}` (${ts.info.type_name()}) is not a supported local type')
        }
    }

    g.local_vars << v
    return v
}

pub fn (mut g Gen) literal_to_constant_expression(typ_ ast.Type, init ast.Expr) ?wasm.ConstExpression {
    typ := ast.mktyp(typ_)
    match init {
        ast.BoolLiteral {
            return wasm.constexpr_value(int(init.val))
        }
        ast.CharLiteral {
            return wasm.constexpr_value(int(init.val.runes()[0]))
        }
        ast.FloatLiteral {
            if typ == ast.f32_type {
                return wasm.constexpr_value(init.val.f32())
            } else if typ == ast.f64_type {
                return wasm.constexpr_value(init.val.f64())
            }
        }
        ast.IntegerLiteral {
            if !typ.is_pure_int() {
                return none
            }
            t := g.get_wasm_type(typ)
            match t {
                .i32_t { return wasm.constexpr_value(init.val.int()) }
                .i64_t { return wasm.constexpr_value(init.val.i64()) }
                else {}
            }
        }
        ast.Ident {
            mut obj := init.obj
            if obj !in [ast.ConstField, ast.GlobalField] {
                obj = init.scope.find(init.name) or { return none }
            }
            match mut obj {
                ast.ConstField {
                    return g.literal_to_constant_expression(typ, obj.expr)
                }
                ast.GlobalField {
                    return g.literal_to_constant_expression(typ, obj.expr)
                }
                else {
                    return none
                }
            }
        }
        else {}
    }

    return none
}

pub fn (mut g Gen) new_global(name string, typ_ ast.Type, init ast.Expr, is_global_mut bool) Global {
    mut typ := ast.mktyp(typ_)
    ts := g.table.sym(typ)

    match ts.info {
        ast.Enum {
            typ = ts.info.typ
        }
        ast.Alias {
            typ = ts.info.parent_type
        }
        else {}
    }

    mut is_mut := false
    is_address := !g.is_pure_type(typ)
    mut init_expr := ?ast.Expr(none)

    cexpr := if cexpr_v := g.literal_to_constant_expression(unpack_literal_int(typ), init) {
        is_mut = is_global_mut
        cexpr_v
    } else {
        // Isn't a literal ...
        if is_address {
            // ... allocate memory and append
            pos, is_init := g.pool.append(init, typ)
            if !is_init {
                // ... AND wait for init in `_vinit`
                init_expr = init
            }
            wasm.constexpr_value(g.data_base + pos)
        } else {
            // ... wait for init in `_vinit`
            init_expr = init
            is_mut = true

            t := g.get_wasm_type(typ)
            wasm.constexpr_value_zero(t)
        }
    }

    mut glbl := Global{
        init: init_expr
        v:    Var{
            name:       name
            typ:        typ
            is_address: is_address
            is_global:  true
            g_idx:      g.mod.new_global(g.dbg_type_name(name, typ), false,
                g.get_wasm_type_int_literal(typ), is_mut, cexpr)
        }
    }

    return glbl
}

// is_pure_type(voidptr) == true
// is_pure_type(&Struct) == false
pub fn (g &Gen) is_pure_type(typ ast.Type) bool {
    if typ.is_pure_int() || typ.is_pure_float() || typ == ast.char_type_idx
        || typ.is_any_kind_of_pointer() || typ.is_bool() {
        return true
    }
    ts := g.table.sym(typ)
    match ts.info {
        ast.Alias {
            return g.is_pure_type(ts.info.parent_type)
        }
        ast.Enum {
            return g.is_pure_type(ts.info.typ)
        }
        else {}
    }

    return false
}

pub fn log2(size int) int {
    return match size {
        1 { 0 }
        2 { 1 }
        4 { 2 }
        8 { 3 }
        else { panic('unreachable') }
    }
}

pub fn (mut g Gen) load(typ ast.Type, offset int) {
    size, align := g.pool.type_size(typ)
    wtyp := g.as_numtype(g.get_wasm_type(typ))

    match size {
        1 { g.func.load8(wtyp, typ.is_signed(), log2(align), offset) }
        2 { g.func.load16(wtyp, typ.is_signed(), log2(align), offset) }
        else { g.func.load(wtyp, log2(align), offset) }
    }
}

pub fn (mut g Gen) store(typ ast.Type, offset int) {
    size, align := g.pool.type_size(typ)
    wtyp := g.as_numtype(g.get_wasm_type(typ))

    match size {
        1 { g.func.store8(wtyp, log2(align), offset) }
        2 { g.func.store16(wtyp, log2(align), offset) }
        else { g.func.store(wtyp, log2(align), offset) }
    }
}

pub fn (mut g Gen) get(v Var) {
    if v.is_global {
        g.func.global_get(v.g_idx)
    } else {
        g.func.local_get(v.idx)
    }

    if v.is_address && g.is_pure_type(v.typ) {
        g.load(v.typ, v.offset)
    } else if v.is_address && v.offset != 0 {
        g.func.i32_const(i32(v.offset))
        g.func.add(.i32_t)
    }
}

pub fn (mut g Gen) mov(to Var, v Var) {
    if !v.is_address || g.is_pure_type(v.typ) {
        g.get(v)
        g.cast(v.typ, to.typ)
        g.set(to)
        return
    }

    size, _ := g.pool.type_size(v.typ)

    if size > 16 {
        g.ref(to)
        g.ref(v)
        g.func.i32_const(i32(size))
        g.func.memory_copy()
        return
    }

    mut sz := size
    mut oz := 0
    for sz > 0 {
        g.ref_ignore_offset(to)
        g.ref_ignore_offset(v)
        if sz - 8 >= 0 {
            g.load(ast.u64_type_idx, v.offset + oz)
            g.store(ast.u64_type_idx, to.offset + oz)
            sz -= 8
            oz += 8
        } else if sz - 4 >= 0 {
            g.load(ast.u32_type_idx, v.offset + oz)
            g.store(ast.u32_type_idx, to.offset + oz)
            sz -= 4
            oz += 4
        } else if sz - 2 >= 0 {
            g.load(ast.u16_type_idx, v.offset + oz)
            g.store(ast.u16_type_idx, to.offset + oz)
            sz -= 2
            oz += 2
        } else if sz - 1 >= 0 {
            g.load(ast.u8_type_idx, v.offset + oz)
            g.store(ast.u8_type_idx, to.offset + oz)
            sz -= 1
            oz += 1
        }
    }
}

pub fn (mut g Gen) set_prepare(v Var) {
    if !v.is_address {
        return
    }

    if g.is_pure_type(v.typ) {
        if v.is_global {
            g.func.global_get(v.g_idx)
        } else {
            g.func.local_get(v.idx)
        }
        return
    }
}

pub fn (mut g Gen) set_set(v Var) {
    if !v.is_address {
        if v.is_global {
            g.func.global_set(v.g_idx)
        } else {
            g.func.local_set(v.idx)
        }
        return
    }

    if g.is_pure_type(v.typ) {
        g.store(v.typ, v.offset)
        return
    }

    from := Var{
        typ:        v.typ
        idx:        g.func.new_local_named(.i32_t, '__tmp<voidptr>')
        is_address: v.is_address
    }

    g.func.local_set(from.idx)
    g.mov(v, from)
}

// set structures with pointer, memcpy.
// set pointers with value, get local, store value.
// set value, set local.
// -- set works with a single value present on the stack beforehand
// -- not optimal for copying stack memory or shuffling structs
// -- use mov instead
pub fn (mut g Gen) set(v Var) {
    if !v.is_address {
        if v.is_global {
            g.func.global_set(v.g_idx)
        } else {
            g.func.local_set(v.idx)
        }
        return
    }

    if g.is_pure_type(v.typ) {
        l := g.new_local('__tmp', v.typ)
        g.func.local_set(l.idx)

        if v.is_global {
            g.func.global_get(v.g_idx)
        } else {
            g.func.local_get(v.idx)
        }

        g.func.local_get(l.idx)

        g.store(v.typ, v.offset)
        return
    }

    from := Var{
        typ:        v.typ
        idx:        g.func.new_local_named(.i32_t, '__tmp<voidptr>')
        is_address: v.is_address
    }

    g.func.local_set(from.idx)
    g.mov(v, from)
}

// to satisfy inline assembly needs
// never used by actual codegen
pub fn (mut g Gen) tee(v Var) {
    assert !v.is_global

    if !v.is_address {
        g.func.local_tee(v.idx)
        return
    }

    if g.is_pure_type(v.typ) {
        l := g.new_local('__tmp', v.typ)
        g.func.local_tee(l.idx) // tee here, leave on stack

        g.func.local_get(v.idx)
        g.func.local_get(l.idx)
        g.store(v.typ, v.offset)
        return
    }

    from := Var{
        typ:        v.typ
        idx:        g.func.new_local_named(.i32_t, '__tmp<voidptr>')
        is_address: v.is_address
    }

    g.func.local_tee(from.idx) // tee here, leave on stack
    g.mov(v, from)
}

pub fn (mut g Gen) ref(v Var) {
    g.ref_ignore_offset(v)

    if v.offset != 0 {
        g.func.i32_const(i32(v.offset))
        g.func.add(.i32_t)
    }
}

pub fn (mut g Gen) ref_ignore_offset(v Var) {
    if !v.is_address {
        panic('unreachable')
    }

    if v.is_global {
        g.func.global_get(v.g_idx)
    } else {
        g.func.local_get(v.idx)
    }
}

// creates a new pointer variable with the offset `offset` and type `typ`
pub fn (mut g Gen) offset(v Var, typ ast.Type, offset int) Var {
    if !v.is_address {
        panic('unreachable')
    }

    nv := Var{
        ...v
        typ:    typ
        offset: v.offset + offset
    }

    return nv
}

pub fn (mut g Gen) zero_fill(v Var, size int) {
    assert size > 0

    // TODO: support coalescing `zero_fill` calls together.
    //       maybe with some kind of context?
    //
    // ```v
    // struct AA {
    //     a bool
    //     b int = 20
    //     c int
    //     d int
    // }
    // ```
    //
    // ```wast
    // (i32.store8
    //  (local.get $0)
    //  (i32.const 0)
    // )
    // (i32.store offset=4
    //  (i32.const 20)
    //  (local.get $0)
    // )                    ;; /- join these together.
    // (i32.store offset=8  ;;-\
    //  (local.get $0)      ;; |
    //  (i32.const 0)       ;; |
    // )                    ;; |
    // (i32.store offset=12 ;; |
    //  (local.get $0)      ;; |
    //  (i32.const 0)       ;; |
    // )                    ;;-/
    // ```

    if size > 16 {
        g.ref(v)
        g.func.i32_const(0)
        g.func.i32_const(i32(size))
        g.func.memory_fill()
        return
    }

    mut sz := size
    mut oz := 0
    for sz > 0 {
        g.ref_ignore_offset(v)
        if sz - 8 >= 0 {
            g.func.i64_const(0)
            g.store(ast.u64_type_idx, v.offset + oz)
            sz -= 8
            oz += 8
        } else if sz - 4 >= 0 {
            g.func.i32_const(0)
            g.store(ast.u32_type_idx, v.offset + oz)
            sz -= 4
            oz += 4
        } else if sz - 2 >= 0 {
            g.func.i32_const(0)
            g.store(ast.u16_type_idx, v.offset + oz)
            sz -= 2
            oz += 2
        } else if sz - 1 >= 0 {
            g.func.i32_const(0)
            g.store(ast.u8_type_idx, v.offset + oz)
            sz -= 1
            oz += 1
        }
    }
}

pub fn (g &Gen) is_param(v Var) bool {
    if v.is_global {
        return false
    }

    return v.idx < g.fn_local_idx_end
}

pub fn (mut g Gen) set_with_multi_expr(init ast.Expr, expected ast.Type, existing_rvars []Var) {
    // misleading name: this doesn't really perform similar to `set_with_expr`
    match init {
        ast.ConcatExpr {
            mut r := 0
            types := g.unpack_type(expected)
            for idx, expr in init.vals {
                typ := types[idx]
                if g.is_param_type(typ) {
                    if rhs := g.get_var_from_expr(expr) {
                        g.mov(existing_rvars[r], rhs)
                    } else {
                        g.set_with_expr(expr, existing_rvars[r])
                    }
                    r++
                } else {
                    g.expr(expr, typ)
                }
            }
        }
        ast.IfExpr {
            g.if_expr(init, expected, existing_rvars)
        }
        ast.MatchExpr {
            g.match_expr(init, expected, existing_rvars)
        }
        ast.CallExpr {
            g.call_expr(init, expected, existing_rvars)
        }
        else {
            if existing_rvars.len > 1 {
                g.w_error('wasm.set_with_multi_expr(): (existing_rvars.len > 1) node: ' +
                    init.type_name())
            }

            if existing_rvars.len == 1 && g.is_param_type(expected) {
                if rhs := g.get_var_from_expr(init) {
                    g.mov(existing_rvars[0], rhs)
                } else {
                    g.set_with_expr(init, existing_rvars[0])
                }
            } else {
                g.expr(init, expected)
            }
        }
    }
}

pub fn (mut g Gen) set_with_expr(init ast.Expr, v Var) {
    match init {
        ast.StructInit {
            size, _ := g.pool.type_size(v.typ)
            ts := g.table.sym(v.typ)
            ts_info := ts.info as ast.Struct
            si := g.pool.type_struct_info(v.typ) or { panic('unreachable') }

            if init.init_fields.len == 0 && !(ts_info.fields.any(it.has_default_expr)) {
                // Struct definition contains no default initialisers
                // AND struct init contains no set values.
                g.zero_fill(v, size)
                return
            }

            for i, f in ts_info.fields {
                field_to_be_set := init.init_fields.map(it.name).contains(f.name)

                if !field_to_be_set {
                    offset := si.offsets[i]
                    offset_var := g.offset(v, f.typ, offset)

                    fsize, _ := g.pool.type_size(f.typ)

                    if f.has_default_expr {
                        g.set_with_expr(f.default_expr, offset_var)
                    } else {
                        g.zero_fill(offset_var, fsize)
                    }
                }
            }

            for f in init.init_fields {
                field := ts.find_field(f.name) or {
                    g.w_error('could not find field `${f.name}` on init')
                }

                offset := si.offsets[field.i]
                offset_var := g.offset(v, f.expected_type, offset)

                g.set_with_expr(f.expr, offset_var)
            }
        }
        ast.StringLiteral {
            val := serialise.eval_escape_codes(init) or { panic('unreachable') }
            str_pos := g.pool.append_string(val)

            if v.typ != ast.string_type {
                // c'str'
                g.set_prepare(v)
                {
                    g.literalint(g.data_base + str_pos, ast.voidptr_type)
                }
                g.set_set(v)
                return
            }

            // init struct
            // fields: str, len
            g.ref(v)
            g.literalint(g.data_base + str_pos, ast.voidptr_type)
            g.store_field(ast.string_type, ast.voidptr_type, 'str')
            g.ref(v)
            g.literalint(val.len, ast.int_type)
            g.store_field(ast.string_type, ast.int_type, 'len')
        }
        ast.CallExpr {
            // `set_with_expr` is never called with a multireturn call expression
            is_pt := g.is_param_type(v.typ)

            g.call_expr(init, v.typ, if is_pt { [v] } else { []Var{} })
            if !is_pt {
                g.set(v)
            }
        }
        ast.ArrayInit {
            if !init.is_fixed {
                g.v_error('wasm backend does not support non fixed arrays yet', init.pos)
            }

            elm_typ := init.elem_type
            elm_size, _ := g.pool.type_size(elm_typ)

            if !init.has_val {
                arr_size, _ := g.pool.type_size(v.typ)

                g.zero_fill(v, arr_size)
                return
            }

            mut voff := g.offset(v, elm_typ, 0) // index zero

            for e in init.exprs {
                g.set_with_expr(e, voff)
                voff = g.offset(voff, elm_typ, elm_size)
            }
        }
        else {
            // impl of set but taken out

            if !v.is_address {
                g.expr(init, v.typ)
                if v.is_global {
                    g.func.global_set(v.g_idx)
                } else {
                    g.func.local_set(v.idx)
                }
                return
            }

            if g.is_pure_type(v.typ) {
                if v.is_global {
                    g.func.global_get(v.g_idx)
                } else {
                    g.func.local_get(v.idx)
                }
                g.expr(init, v.typ)
                g.store(v.typ, v.offset)
                return
            }

            if var := g.get_var_from_expr(init) {
                g.mov(v, var)
                return
            }

            from := Var{
                typ:        v.typ
                idx:        g.func.new_local_named(.i32_t, '__tmp<voidptr>')
                is_address: v.is_address // true
            }

            g.expr(init, v.typ)
            g.func.local_set(from.idx)
            g.mov(v, from)
        }
    }
}

pub fn calc_padding(value int, alignment int) int {
    if alignment == 0 {
        return value
    }
    return (alignment - value % alignment) % alignment
}

pub fn calc_align(value int, alignment int) int {
    if alignment == 0 {
        return value
    }
    return (value + alignment - 1) / alignment * alignment
}

pub fn (mut g Gen) make_vinit() {
    g.func = g.mod.new_function('_vinit', [], [])
    func_start := g.func.patch_pos()
    {
        for mod_name in g.table.modules {
            if mod_name == 'v.reflection' {
                g.w_error('the wasm backend does not implement `v.reflection` yet')
            }
            init_fn_name := if mod_name != 'builtin' { '${mod_name}.init' } else { 'init' }
            if _ := g.table.find_fn(init_fn_name) {
                g.func.call(init_fn_name)
            }
            cleanup_fn_name := if mod_name != 'builtin' { '${mod_name}.cleanup' } else { 'cleanup' }
            if _ := g.table.find_fn(cleanup_fn_name) {
                g.func.call(cleanup_fn_name)
            }
        }
        for _, gv in g.global_vars {
            if init := gv.init {
                g.set_with_expr(init, gv.v)
            }
        }
        g.bare_function_frame(func_start)
    }
    g.mod.commit(g.func, false)
    g.bare_function_end()
}

pub fn (mut g Gen) housekeeping() {
    g.make_vinit()

    heap_base := calc_align(g.data_base + g.pool.buf.len, 16) // 16?
    page_boundary := calc_align(g.data_base + g.pool.buf.len, 64 * 1024)
    preallocated_pages := page_boundary / (64 * 1024)

    if g.pref.is_verbose {
        eprintln('housekeeping(): acceptable addresses are > 1024')
        eprintln('housekeeping(): stack top: ${g.stack_top}, data_base: ${g.data_base} (size: ${g.pool.buf.len}), heap_base: ${heap_base}')
        eprintln('housekeeping(): preallocated pages: ${preallocated_pages}')
    }

    if sp := g.sp_global {
        g.mod.assign_global_init(sp, wasm.constexpr_value(g.stack_top))
    }
    if g.sp_global != none || g.pool.buf.len > 0 {
        g.mod.assign_memory('memory', true, u32(preallocated_pages), none)
        if g.pool.buf.len > 0 {
            mut buf := g.pool.buf.clone()

            for reloc in g.pool.relocs {
                binary.little_endian_put_u32_at(mut buf, u32(g.data_base + reloc.offset), reloc.pos)
            }
            g.mod.new_data_segment(none, g.data_base, buf)
        }
    }
    if hp := g.heap_base {
        g.mod.assign_global_init(hp, wasm.constexpr_value(heap_base))
    }

    if g.pref.os == .wasi && !g.pref.is_shared {
        mut fn_start := g.mod.new_function('_start', [], [])
        {
            fn_start.call('_vinit')
            fn_start.call('main.main')
        }
        g.mod.commit(fn_start, true)
    } else {
        g.mod.assign_start('_vinit')
    }
}