module binary

struct EncodeState {
mut:
    b []u8
    // pre-allocated buffers
    b2         []u8 = [u8(0), 0]
    b4         []u8 = [u8(0), 0, 0, 0]
    b8         []u8 = [u8(0), 0, 0, 0, 0, 0, 0, 0]
    big_endian bool
}

@[params]
pub struct EncodeConfig {
pub mut:
    buffer_len int = 1024
    big_endian bool // use big endian encoding the data
}

// encode_binary encode a T type data into u8 array.
// for encoding struct, you can use `@[serialize: '-']` to skip field.
// Note: `shared` fields in struct will be skipped.
pub fn encode_binary[T](obj T, config EncodeConfig) ![]u8 {
    mut s := EncodeState{
        b:          []u8{cap: config.buffer_len}
        big_endian: config.big_endian
    }
    $if T is $array {
        encode_array(mut s, obj)!
    } $else $if T is $string {
        encode_string(mut s, obj)!
    } $else $if T is $struct {
        encode_struct(mut s, obj)!
    } $else $if T is $map {
        encode_map(mut s, obj)!
    } $else {
        encode_primitive(mut s, obj)!
    }
    return s.b
}

@[inline]
fn normalize_attr_arg(value string) string {
    mut normalized := value.trim_space()
    if normalized.len > 1 {
        if (normalized[0] == `'` && normalized[normalized.len - 1] == `'`)
            || (normalized[0] == `"` && normalized[normalized.len - 1] == `"`) {
            normalized = normalized[1..normalized.len - 1]
        }
    }
    return normalized
}

fn encode_struct[T](mut s EncodeState, obj T) ! {
    $for field in T.fields {
        mut is_skip := false
        for attr in field.attrs {
            f := attr.split_any(':')
            if f.len == 2 {
                match f[0].trim_space() {
                    'serialize' {
                        // @[serialize:'-']
                        if normalize_attr_arg(f[1]) == '-' {
                            is_skip = true
                        }
                    }
                    else {}
                }
            }
        }
        if !is_skip {
            // TODO: support shared field in struct, currently skip.
            $if field.typ !is $shared {
                value := obj.$(field.name)
                $if field.typ is $array {
                    encode_array(mut s, value)!
                } $else $if field.typ is $string {
                    encode_string(mut s, value)!
                } $else $if field.typ is $struct {
                    encode_struct(mut s, value)!
                } $else $if field.typ is $map {
                    encode_map(mut s, value)!
                } $else {
                    encode_primitive(mut s, value)!
                }
            }
        }
    }
}

// help unions for bypass `-cstrict`/ `-Wstrict-aliasing` check.
union U32_F32 {
    u u32
    f f32
}

union U64_F64 {
    u u64
    f f64
}

fn encode_primitive[T](mut s EncodeState, value T) ! {
    $if T is int {
        // NOTE: `int` always use 64bit
        s.put_u64(u64(value))
    } $else $if T is u8 {
        s.put_u8(u8(value))
    } $else $if T is u16 {
        s.put_u16(u16(value))
    } $else $if T is u32 {
        s.put_u32(u32(value))
    } $else $if T is u64 {
        s.put_u64(u64(value))
    } $else $if T is i8 {
        s.put_u8(u8(value))
    } $else $if T is i16 {
        s.put_u16(u16(value))
    } $else $if T is i32 {
        s.put_u32(u32(value))
    } $else $if T is i64 {
        s.put_u64(u64(value))
    } $else $if T is f32 {
        unsafe { s.put_u32(U32_F32{ f: value }.u) }
    } $else $if T is f64 {
        unsafe { s.put_u64(U64_F64{ f: value }.u) }
    } $else $if T is bool {
        s.put_u8(u8(value))
    } $else $if T is rune {
        s.put_u32(u32(value))
    } $else $if T is isize {
        if sizeof(isize) == 4 {
            s.put_u32(u32(value))
        } else {
            s.put_u64(u64(value))
        }
    } $else $if T is usize {
        if sizeof(usize) == 4 {
            s.put_u32(u32(value))
        } else {
            s.put_u64(u64(value))
        }
    } $else $if T is voidptr {
        s.put_u64(u64(value))
    } $else {
        // TODO: `any` type support?
        return error('${@FN}(): unsupported type ${typeof(value).name}')
    }
}

fn encode_array[T](mut s EncodeState, arr []T) ! {
    s.put_u64(u64(arr.len))

    $if T is u8 {
        // optimization for `[]u8`
        s.b << arr
    } $else {
        for element in arr {
            $if T is $string {
                encode_string(mut s, element)!
            } $else $if T is $struct {
                encode_struct(mut s, element)!
            } $else $if T is $map {
                encode_map(mut s, element)!
            } $else {
                encode_primitive(mut s, element)!
            }
        }
    }
}

fn encode_string(mut s EncodeState, str string) ! {
    s.put_u64(u64(str.len))
    s.b << str.bytes()
}

fn encode_map[K, V](mut s EncodeState, m map[K]V) ! {
    s.put_u64(u64(m.len))

    for k, v in m {
        // encode key first
        // Maps can have keys of type string, rune, integer, float or voidptr.
        $if K is $string {
            encode_string(mut s, k)!
        } $else {
            encode_primitive(mut s, k)!
        }

        // encode value
        $if V is $string {
            encode_string(mut s, v)!
        } $else $if V is $struct {
            encode_struct(mut s, v)!
        } $else $if V is $map {
            encode_map(mut s, v)!
        } $else {
            encode_primitive(mut s, v)!
        }
    }
}

struct DecodeState {
mut:
    b          []u8
    b2         []u8 = [u8(0), 0]
    b4         []u8 = [u8(0), 0, 0, 0]
    b8         []u8 = [u8(0), 0, 0, 0, 0, 0, 0, 0]
    offset     int
    big_endian bool
}

@[params]
pub struct DecodeConfig {
pub mut:
    buffer_len int = 1024
    big_endian bool // use big endian decode the data
}

// decode_binary decode a u8 array into T type data.
// for decoding struct, you can use `@[serialize: '-']` to skip field.
// Note: `shared` fields in struct will be skipped.
pub fn decode_binary[T](b []u8, config DecodeConfig) !T {
    mut s := DecodeState{
        b:          b
        big_endian: config.big_endian
    }
    $if T is $array {
        return decode_array(mut s, T{})!
    } $else $if T is $string {
        return decode_string(mut s)!
    } $else $if T is $struct {
        return decode_struct(mut s, T{})!
    } $else $if T is $map {
        return decode_map(mut s, T{})!
    } $else {
        return decode_primitive(mut s, unsafe { T(0) })!
    }
}

fn decode_struct[T](mut s DecodeState, _ T) !T {
    mut obj := T{}

    $for field in T.fields {
        mut is_skip := false
        for attr in field.attrs {
            f := attr.split_any(':')
            if f.len == 2 {
                match f[0].trim_space() {
                    'serialize' {
                        // @[serialize:'-']
                        if normalize_attr_arg(f[1]) == '-' {
                            is_skip = true
                        }
                    }
                    else {}
                }
            }
        }
        if !is_skip {
            // TODO: support shared field in struct, currently skip.
            $if field.typ !is $shared {
                $if field.typ is $array {
                    obj.$(field.name) = decode_array(mut s, obj.$(field.name))!
                } $else $if field.typ is $string {
                    obj.$(field.name) = decode_string(mut s)!
                } $else $if field.typ is $struct {
                    obj.$(field.name) = decode_struct(mut s, obj.$(field.name))!
                } $else $if field.typ is $map {
                    obj.$(field.name) = decode_map(mut s, obj.$(field.name))!
                } $else {
                    obj.$(field.name) = decode_primitive(mut s, obj.$(field.name))!
                }
            }
        }
    }
    return obj
}

fn decode_primitive[T](mut s DecodeState, value T) !T {
    $if T is int {
        // NOTE: int always use 64bit
        return T(s.get_u64()!)
    } $else $if T is u8 {
        return T(s.get_u8()!)
    } $else $if T is u16 {
        return T(s.get_u16()!)
    } $else $if T is u32 {
        return T(s.get_u32()!)
    } $else $if T is u64 {
        return T(s.get_u64()!)
    } $else $if T is i8 {
        return T(s.get_u8()!)
    } $else $if T is i16 {
        return T(s.get_u16()!)
    } $else $if T is i32 {
        return T(s.get_u32()!)
    } $else $if T is i64 {
        return T(s.get_u64()!)
    } $else $if T is f32 {
        v := s.get_u32()!
        return unsafe {
            U32_F32{
                u: v
            }.f
        }
    } $else $if T is f64 {
        v := s.get_u64()!
        return unsafe {
            U64_F64{
                u: v
            }.f
        }
    } $else $if T is bool {
        return s.get_u8()! != 0
    } $else $if T is rune {
        return T(s.get_u32()!)
    } $else $if T is isize {
        if sizeof(isize) == 4 {
            return T(s.get_u32()!)
        } else {
            return T(s.get_u64()!)
        }
    } $else $if T is usize {
        if sizeof(usize) == 4 {
            return T(s.get_u32()!)
        } else {
            return T(s.get_u64()!)
        }
    } $else $if T is voidptr {
        return T(s.get_u64()!)
    } $else {
        // TODO: `any` type support?
        return error('${@FN}(): unsupported type ${typeof(value).name}')
    }
    return error('${@FN}(): impossible error')
}

fn decode_array[T](mut s DecodeState, _ []T) ![]T {
    len := int(s.get_u64()!)
    if len <= 0 || s.offset + len > s.b.len {
        return error('${@FN}(): invalid array length decode from stream')
    }
    mut arr := []T{cap: len}
    $if T is u8 {
        // optimization for `[]u8`
        arr << s.b[s.offset..s.offset + len]
        s.offset += len
    } $else {
        for _ in 0 .. len {
            if s.offset >= s.b.len {
                return error('${@FN}(): unexpected end of data')
            }
            $if T is $array {
                arr << decode_array(mut s, T{})!
            } $else $if T is $string {
                arr << decode_string(mut s)!
            } $else $if T is $struct {
                arr << decode_struct(mut s, T{})!
            } $else $if T is $map {
                arr << decode_map(mut s, T{})!
            } $else {
                arr << decode_primitive(mut s, unsafe { T(0) })!
            }
        }
    }
    return arr
}

fn decode_string(mut s DecodeState) !string {
    len := int(s.get_u64()!)
    if len <= 0 || s.offset + len > s.b.len {
        return error('${@FN}(): invalid string length decode from stream')
    }
    str := unsafe { s.b[s.offset..s.offset + len].bytestr() }
    s.offset += len
    return str
}

// `Any` is a sum type that lists the possible types to be decoded and used.
type Any = int
    | bool
    | f64
    | f32
    | i64
    | i32
    | i16
    | i8
    | map[string]Any
    | map[int]Any
    | string
    | u64
    | u32
    | u16
    | u8
    | rune
    | isize
    | usize
    | []Any

fn decode_map[K, V](mut s DecodeState, _ map[K]V) !map[K]V {
    len := int(s.get_u64()!)
    if len <= 0 || s.offset + len > s.b.len {
        return error('${@FN}(): invalid map length decode from stream')
    }

    mut m := map[K]V{}

    for _ in 0 .. len {
        // decode key first
        // Maps can have keys of type string, rune, integer, float or voidptr.
        mut k := Any(0)
        $if K is $string {
            k = decode_string(mut s)!
        } $else {
            k = decode_primitive(mut s, unsafe { K(0) })!
        }

        // decode value
        $if V is $struct {
            v := decode_struct(mut s, V{})!
            m[k as K] = v
        } $else $if V is $map {
            v := decode_map(mut s, V{})!
            m[k as K] = v
        } $else $if V is $string {
            v := decode_string(mut s)!
            m[k as K] = v
        } $else {
            v := decode_primitive(mut s, unsafe { V(0) })!
            m[k as K] = v
        }
    }
    return m
}

@[inline]
fn (mut s DecodeState) get_u64() !u64 {
    if s.offset + 8 > s.b.len {
        return error('${@FN}(): bytes length is not enough for u64')
    }
    defer {
        s.offset += 8
    }
    if s.big_endian {
        return big_endian_u64_at(s.b, s.offset)
    } else {
        return little_endian_u64_at(s.b, s.offset)
    }
}

@[inline]
fn (mut s DecodeState) get_u32() !u32 {
    if s.offset + 4 > s.b.len {
        return error('${@FN}(): bytes length is not enough for u32')
    }
    defer {
        s.offset += 4
    }
    if s.big_endian {
        return big_endian_u32_at(s.b, s.offset)
    } else {
        return little_endian_u32_at(s.b, s.offset)
    }
}

@[inline]
fn (mut s DecodeState) get_u16() !u16 {
    if s.offset + 2 > s.b.len {
        return error('${@FN}(): bytes length is not enough for u16')
    }
    defer {
        s.offset += 2
    }
    if s.big_endian {
        return big_endian_u16_at(s.b, s.offset)
    } else {
        return little_endian_u16_at(s.b, s.offset)
    }
}

@[inline]
fn (mut s DecodeState) get_u8() !u8 {
    if s.offset + 1 > s.b.len {
        return error('${@FN}(): bytes length is not enough for u8')
    }
    defer {
        s.offset += 1
    }
    return s.b[s.offset]
}

@[inline]
fn (mut s EncodeState) put_u64(value u64) {
    if s.big_endian {
        big_endian_put_u64(mut s.b8, value)
    } else {
        little_endian_put_u64(mut s.b8, value)
    }
    s.b << s.b8
}

@[inline]
fn (mut s EncodeState) put_u32(value u32) {
    if s.big_endian {
        big_endian_put_u32(mut s.b4, value)
    } else {
        little_endian_put_u32(mut s.b4, value)
    }
    s.b << s.b4
}

@[inline]
fn (mut s EncodeState) put_u16(value u16) {
    if s.big_endian {
        big_endian_put_u16(mut s.b2, value)
    } else {
        little_endian_put_u16(mut s.b2, value)
    }
    s.b << s.b2
}

@[inline]
fn (mut s EncodeState) put_u8(value u8) {
    s.b << value
}