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

import v.ast
// import v.eval
import math.bits
import strconv

@[noinit]
pub struct Pool {
mut:
    table &ast.Table
    // eval    eval.Eval
    structs map[ast.Type]StructInfo
    strings []StringInfo // string intern
pub:
    null_terminated bool
    intern_strings  bool
    store_relocs    bool
pub mut:
    buf               []u8
    relocs            []Reloc
    highest_alignment int
}

struct StringInfo {
    pos int
    len int
}

pub struct StructInfo {
pub mut:
    offsets []int
}

pub struct Reloc {
pub:
    pos    int
    offset int
}

pub fn (mut p Pool) type_struct_info(typ ast.Type) ?StructInfo {
    ts := p.table.sym(typ)

    if ts.info !is ast.Struct {
        return none
    }

    if typ.idx() in p.structs {
        return p.structs[typ.idx()]
    }

    // will cache inside `p.structs`
    p.type_size(typ)
    return p.structs[typ.idx()]
}

pub fn (mut p Pool) type_size(typ ast.Type) (int, int) {
    if typ.nr_muls() > 0 {
        return p.table.pointer_size, p.table.pointer_size
    }

    ts := p.table.sym(typ)
    if ts.size != -1 && typ.idx() in p.structs {
        return ts.size, ts.align
    }

    if ts.info is ast.Enum {
        return p.table.type_size(ts.info.typ)
    }

    if ts.info !is ast.Struct {
        return p.table.type_size(typ)
    }

    ti := ts.info as ast.Struct

    // code borrowed from native, inserted in wasm, and now here!

    mut strc := StructInfo{}
    mut size := 0
    mut align := 1
    for f in ti.fields {
        f_size, f_align := p.type_size(f.typ)
        if f_size == 0 {
            strc.offsets << 0
            continue
        }
        padding := (f_align - size % f_align) % f_align
        strc.offsets << size + padding
        size += f_size + padding
        if f_align > align {
            align = f_align
        }
    }
    size = (size + align - 1) / align * align
    p.structs[typ.idx()] = strc

    mut ts_ := p.table.sym(typ)
    ts_.size = size
    ts_.align = align

    return size, align
}

@[params]
pub struct PoolOpts {
pub:
    null_terminated bool = true
    intern_strings  bool = true
    store_relocs    bool = true
}

pub fn new_pool(table &ast.Table, opts PoolOpts) Pool {
    return Pool{
        table:           table
        null_terminated: opts.null_terminated
        intern_strings:  opts.intern_strings
        store_relocs:    opts.store_relocs
    }
}

fn (mut p Pool) zero_fill(size int) {
    // TODO: eventually support a way to utilise a BSS section

    for i := 0; i < size; i++ {
        p.buf << 0
    }
}

fn (mut p Pool) alignment(align int) int {
    if align > p.highest_alignment {
        p.highest_alignment = align
    }
    padding := (align - p.buf.len % align) % align
    p.zero_fill(padding)
    pos := p.buf.len
    return pos
}

/*
fn (mut p Pool) append_struct(init ast.StructInit) ?int {
    old_len := p.buf.len

    size, align := p.type_size(v.typ)
    ts := g.table.sym(v.typ)
    ts_info := ts.info as ast.Struct

    pos := p.alignment(align)
    
    if init.fields.len == 0 && !(ts_info.fields.any(it.has_default_expr)) {
        for i := 0 ; i < size ; i++ {
            p.buf << 0
        }
        return pos
    }

    /* for i, f in ts_info.fields {
        field_to_be_set := init.fields.map(it.name).filter(f.name)
    } */

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

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

            fsize, _ := g.get_type_size_align(f.typ)
            
            if f.has_default_expr {
                g.expr(f.default_expr, f.typ)
                g.set(offset_var)
            } else {
                g.zero_fill(offset_var, fsize)
            }                        
        }
    }

    for f in init.fields {
        field := ts.find_field(f.name) or {
            g.w_error('could not find field `${f.name}` on init')
        }
        
        offset := g.structs[v.typ.idx()].offsets[field.i]
        offset_var := g.offset(v, f.expected_type, offset)

        g.expr(f.expr, f.expected_type)
        g.set(offset_var)
    } */

    return pos
}*/

pub fn eval_escape_codes_raw(str string) !string {
    mut buffer := []u8{}

    mut i := 0
    for i < str.len {
        if str[i] != `\\` {
            buffer << str[i]
            i++
            continue
        }

        // skip \
        i++
        match str[i] {
            `\\`, `'`, `"` {
                buffer << str[i]
                i++
            }
            `a`, `b`, `f` {
                buffer << str[i] - u8(90)
                i++
            }
            `n` {
                buffer << `\n`
                i++
            }
            `r` {
                buffer << `\r`
                i++
            }
            `t` {
                buffer << `\t`
                i++
            }
            `u` {
                i++
                utf8 := strconv.parse_int(str[i..i + 4], 16, 16) or {
                    return error('invalid \\u escape code (${str[i..i + 4]})')
                }
                i += 4
                buffer << u8(utf8)
                buffer << u8(utf8 >> 8)
            }
            `v` {
                buffer << `\v`
                i++
            }
            `x` {
                i++
                c := strconv.parse_int(str[i..i + 2], 16, 8) or {
                    return error('invalid \\x escape code (${str[i..i + 2]})')
                }
                i += 2
                buffer << u8(c)
            }
            `0`...`7` {
                c := strconv.parse_int(str[i..i + 3], 8, 8) or {
                    return error('invalid escape code \\${str[i..i + 3]}')
                }
                i += 3
                buffer << u8(c)
            }
            else {
                return error('invalid escape code \\${str[i]}')
            }
        }
    }

    return buffer.bytestr()
}

pub fn eval_escape_codes(str_lit ast.StringLiteral) !string {
    if str_lit.is_raw {
        return str_lit.val
    }

    return eval_escape_codes_raw(str_lit.val)
}

pub fn (mut p Pool) append_string(val string) int {
    data := val.bytes()

    if p.intern_strings {
        for str in p.strings {
            if data.len > str.len || (p.null_terminated && data.len != str.len) {
                continue
            }

            // TODO: aggressive string interning if `p.null_terminated`
            if p.buf[str.pos..str.pos + data.len] == data {
                return str.pos
            }
        }
    }

    pos := p.buf.len
    p.buf << data
    if p.null_terminated {
        p.buf << 0
    }

    p.strings << StringInfo{
        pos: pos
        len: data.len
    }

    return pos
}

pub fn (mut p Pool) append(init ast.Expr, typ ast.Type) (int, bool) {
    match init {
        ast.BoolLiteral {
            pos := p.buf.len
            p.buf << u8(init.val)
            return pos, true
        }
        ast.FloatLiteral {
            assert typ.is_pure_float()

            mut pos := 0
            if typ == ast.f32_type {
                pos = p.alignment(4)
                p.u32(bits.f32_bits(init.val.f32()))
            } else {
                pos = p.alignment(8)
                p.u64(bits.f64_bits(init.val.f64()))
            }

            return pos, true
        }
        ast.IntegerLiteral {
            assert typ.is_pure_int()

            size, align := p.table.type_size(typ)
            pos := p.alignment(align)

            match size {
                1 {
                    p.u8(u8(init.val.i8()))
                }
                2 {
                    p.u16(u16(init.val.i16()))
                }
                4 {
                    p.u32(u32(init.val.int()))
                }
                8 {
                    p.u64(u64(init.val.i64()))
                }
                else {}
            }

            return pos, true
        }
        ast.CharLiteral {
            // 3 extra bytes for improved program correctness, thank me later
            rne := u32(eval_escape_codes_raw(init.val) or { panic('Pool.append: ${err}') }.runes()[0])
            pos := p.alignment(4)
            p.u32(rne)

            return pos, true
        }
        ast.StringLiteral {
            val := eval_escape_codes(init) or { panic('Pool.append: ${err}') }
            str_pos := p.append_string(val)

            if typ != ast.string_type {
                // c'str'
                return str_pos, true
            }

            _, align := p.type_size(ast.string_type)
            tss := p.table.sym(ast.string_type).info as ast.Struct
            pos := p.alignment(align)

            for field in tss.fields {
                match field.name {
                    'str' {
                        p.ptr(str_pos)
                    }
                    'len' {
                        p.u32(u32(val.len))
                    }
                    'is_lit' {
                        p.u32(1)
                    }
                    else {
                        panic('ast.string: field `${field.name}` is unknown')
                    }
                }
            }

            return pos, true
        }
        else {
            size, align := p.type_size(typ)
            pos := p.alignment(align)
            p.zero_fill(size)
            return pos, false
        }
    }
}

fn (mut p Pool) u64(v u64) {
    p.buf << u8(v)
    p.buf << u8(v >> u64(8))
    p.buf << u8(v >> u64(16))
    p.buf << u8(v >> u64(24))
    p.buf << u8(v >> u64(32))
    p.buf << u8(v >> u64(40))
    p.buf << u8(v >> u64(48))
    p.buf << u8(v >> u64(56))
}

fn (mut p Pool) u32(v u32) {
    p.buf << u8(v)
    p.buf << u8(v >> u32(8))
    p.buf << u8(v >> u32(16))
    p.buf << u8(v >> u32(24))
}

fn (mut p Pool) u16(v u16) {
    p.buf << u8(v)
    p.buf << u8(v >> u32(8))
}

fn (mut p Pool) u8(v u8) {
    p.buf << v
}

fn (mut p Pool) ptr(offset int) int {
    assert p.table.pointer_size in [1, 2, 4, 8]
    pos := p.buf.len // p.alignment(p.table.pointer_size)

    if p.store_relocs {
        p.relocs << Reloc{
            pos:    pos
            offset: offset
        }
    }

    match p.table.pointer_size {
        1 {
            p.u8(u8(offset))
        }
        2 {
            p.u16(u16(offset))
        }
        4 {
            p.u32(u32(offset))
        }
        8 {
            p.u64(u64(offset))
        }
        else {}
    }

    return pos
}