module json2

import time

// EncoderOptions provides a list of options for encoding
@[params]
pub struct EncoderOptions {
pub:
	prettify       bool
	indent_string  string = '    '
	newline_string string = '\n'

	enum_as_int bool

	escape_unicode bool
}

struct Encoder {
	EncoderOptions
mut:
	level  int
	prefix string

	output []u8 = []u8{cap: 2048}
}

// encode is a generic function that encodes a type into a JSON string.
pub fn encode[T](val T, config EncoderOptions) string {
	mut encoder := Encoder{
		EncoderOptions: config
	}

	encoder.encode_value[T](val)

	return encoder.output.bytestr()
}

fn (mut encoder Encoder) encode_value[T](val T) {
	$if T is $interface {
		encoder.encode_null()
	} $else $if T.unaliased_typ is voidptr {
		encoder.encode_null()
	} $else $if T.unaliased_typ is string {
		encoder.encode_string(string(val))
	} $else $if T.unaliased_typ is bool {
		encoder.encode_boolean(bool(val))
	} $else $if T.unaliased_typ is u8 {
		encoder.encode_number(u8(val))
	} $else $if T.unaliased_typ is u16 {
		encoder.encode_number(u16(val))
	} $else $if T.unaliased_typ is u32 {
		encoder.encode_number(u32(val))
	} $else $if T.unaliased_typ is u64 {
		encoder.encode_number(u64(val))
	} $else $if T.unaliased_typ is i8 {
		encoder.encode_number(i8(val))
	} $else $if T.unaliased_typ is i16 {
		encoder.encode_number(i16(val))
	} $else $if T.unaliased_typ is int || T.unaliased_typ is i32 {
		encoder.encode_number(i32(val))
	} $else $if T.unaliased_typ is i64 {
		encoder.encode_number(i64(val))
	} $else $if T.unaliased_typ is usize {
		encoder.encode_number(usize(val))
	} $else $if T.unaliased_typ is isize {
		encoder.encode_number(isize(val))
	} $else $if T.unaliased_typ is f32 {
		encoder.encode_number(f32(val))
	} $else $if T.unaliased_typ is f64 {
		encoder.encode_number(f64(val))
	} $else $if T.unaliased_typ is voidptr {
		encoder.encode_number(0)
	} $else $if T is $pointer {
		if voidptr(val) == unsafe { nil } {
			encoder.encode_null()
		} else {
			encoder.encode_value(*val)
		}
	} $else $if T.unaliased_typ is $array_fixed {
		encoder.output << `[`
		for i in 0 .. val.len {
			encoder.encode_value(val[i])
			if i < val.len - 1 {
				encoder.output << `,`
			}
		}
		encoder.output << `]`
	} $else $if T.unaliased_typ is $array {
		encoder.encode_array(val)
	} $else $if T.unaliased_typ is $map {
		encoder.output << `{`
		if encoder.prettify {
			encoder.increment_level()
			encoder.add_indent()
		}
		mut mi := 0
		for key, value in val {
			encoder.encode_string('${key}')
			encoder.output << `:`
			if encoder.prettify {
				encoder.output << ` `
			}
			encoder.encode_value(value)
			if mi < val.len - 1 {
				encoder.output << `,`
				if encoder.prettify {
					encoder.add_indent()
				}
			} else {
				if encoder.prettify {
					encoder.decrement_level()
					encoder.add_indent()
				}
			}
			mi++
		}
		encoder.output << `}`
	} $else $if T.unaliased_typ is $enum {
		if encoder.enum_as_int || enum_uses_json_as_number[T]() {
			encoder.encode_number(int(val))
		} else {
			mut enum_val := 'unknown enum value'
			$for member in T.values {
				if member.value == val {
					enum_val = member.name
					for attr in member.attrs {
						if json_attr := json_attr_value(attr) {
							enum_val = json_attr
						}
					}
				}
			}
			encoder.output << `"`
			unsafe { encoder.output.push_many(enum_val.str, enum_val.len) }
			encoder.output << `"`
		}
	} $else $if T.unaliased_typ is $sumtype {
		encoder.encode_sumtype[T](val)
	} $else $if T is JsonEncoder { // uses T, because alias could be implementing JsonEncoder, while the base type does not
		integer_val := val.to_json()
		unsafe { encoder.output.push_many(integer_val.str, integer_val.len) }
	} $else $if T is Encodable { // uses T, because alias could be implementing JsonEncoder, while the base type does not
		integer_val := val.json_str()
		unsafe { encoder.output.push_many(integer_val.str, integer_val.len) }
	} $else $if T.unaliased_typ is $struct {
		unsafe {
			$for field in T.fields {
				$if field.is_embed {
					encoder.encode_struct_with_embeds(val)
					return
				}
			}
			encoder.output << `{`
			is_first := encoder.encode_struct_fields[T](val, true, [], '')
			if encoder.prettify && !is_first {
				encoder.decrement_level()
				encoder.add_indent()
			}
			encoder.output << `}`
		}
	}
}

fn (mut encoder Encoder) encode_string(val string) {
	encoder.output << `"`
	mut buffer_start := 0
	mut buffer_end := 0
	for buffer_end < val.len {
		character := val[buffer_end]
		match character {
			`"`, `\\` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << character
			}
			`\b` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << `b`
			}
			`\n` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << `n`
			}
			`\f` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << `f`
			}
			`\t` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << `t`
			}
			`\r` {
				unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }
				buffer_end++
				buffer_start = buffer_end

				encoder.output << `\\`
				encoder.output << `r`
			}
			else {
				if character < 0x20 { // control characters
					unsafe {
						encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start)
					}
					buffer_end++
					buffer_start = buffer_end

					encoder.output << `\\`
					encoder.output << `u`

					hex_string := '${character:04x}'

					unsafe { encoder.output.push_many(hex_string.str, 4) }

					continue
				}
				if encoder.escape_unicode {
					if character >= 0b1111_0000 { // four bytes
						unsafe {
							encoder.output.push_many(val.str + buffer_start,
								buffer_end - buffer_start)
						}
						unicode_point_low := val[buffer_end..buffer_end + 4].bytes().byterune() or {
							0
						} - 0x10000

						hex_string := '\\u${0xD800 + ((unicode_point_low >> 10) & 0x3FF):04X}\\u${
							0xDC00 + (unicode_point_low & 0x3FF):04x}'

						buffer_end += 4
						buffer_start = buffer_end

						unsafe { encoder.output.push_many(hex_string.str, 12) }

						continue
					} else if character >= 0b1110_0000 { // three bytes
						unsafe {
							encoder.output.push_many(val.str + buffer_start,
								buffer_end - buffer_start)
						}
						hex_string := '\\u${val[buffer_end..buffer_end + 3].bytes().byterune() or {
							0
						}:04x}'

						buffer_end += 3
						buffer_start = buffer_end

						unsafe { encoder.output.push_many(hex_string.str, 6) }

						continue
					} else if character >= 0b1100_0000 { // two bytes
						unsafe {
							encoder.output.push_many(val.str + buffer_start,
								buffer_end - buffer_start)
						}
						hex_string := '\\u${val[buffer_end..buffer_end + 2].bytes().byterune() or {
							0
						}:04x}'

						buffer_end += 2
						buffer_start = buffer_end

						unsafe { encoder.output.push_many(hex_string.str, 6) }

						continue
					}
				}

				buffer_end++
			}
		}
	}
	unsafe { encoder.output.push_many(val.str + buffer_start, buffer_end - buffer_start) }

	encoder.output << `"`
}

fn (mut encoder Encoder) encode_boolean(val bool) {
	if val {
		unsafe { encoder.output.push_many(true_string.str, true_string.len) }
	} else {
		unsafe { encoder.output.push_many(false_string.str, false_string.len) }
	}
}

fn (mut encoder Encoder) encode_number[T](val T) {
	mut integer_val := ''
	$if T is u8 {
		integer_val = u8(val).str()
	} $else $if T is u16 {
		integer_val = u16(val).str()
	} $else $if T is u32 {
		integer_val = u32(val).str()
	} $else $if T is u64 {
		integer_val = u64(val).str()
	} $else $if T is i8 {
		integer_val = i8(val).str()
	} $else $if T is i16 {
		integer_val = i16(val).str()
	} $else $if T is int || T is i32 {
		integer_val = i32(val).str()
	} $else $if T is i64 {
		integer_val = i64(val).str()
	} $else $if T is usize {
		integer_val = usize(val).str()
	} $else $if T is isize {
		integer_val = isize(val).str()
	} $else $if T is f32 {
		integer_val = f32(val).str()
	} $else $if T is f64 {
		integer_val = f64(val).str()
	}
	$if T is $float {
		if integer_val.len > 2 && integer_val[integer_val.len - 2] == `.`
			&& integer_val[integer_val.len - 1] == `0` { // ends in .0
			// `2.0` = > `2`
			// but skip float in scientific notation, `1e+10`
			unsafe {
				integer_val.len -= 2
			}
		}
	}
	unsafe { encoder.output.push_many(integer_val.str, integer_val.len) }
}

@[markused]
fn (mut encoder Encoder) encode_null() {
	unsafe { encoder.output.push_many(null_string.str, null_string.len) }
}

fn (mut encoder Encoder) encode_array[T](val T) {
	encoder.output << `[`
	if encoder.prettify {
		encoder.increment_level()
		encoder.add_indent()
	}

	for i, item in val {
		$if T is $pointer {
			if voidptr(item) == unsafe { nil } {
				encoder.encode_null()
			} else {
				unsafe { encoder.encode_pointer_array_item(item) }
			}
		} $else {
			encoder.encode_value(item)
		}
		if i < val.len - 1 {
			encoder.output << `,`
			if encoder.prettify {
				encoder.add_indent()
			}
		} else {
			if encoder.prettify {
				encoder.decrement_level()
				encoder.add_indent()
			}
		}
	}

	encoder.output << `]`
}

@[unsafe]
fn (mut encoder Encoder) encode_pointer_array_item[T](item T) {
	encoder.encode_value(*item)
}

fn (mut encoder Encoder) encode_map[K, T](val map[K]T) {
	encoder.output << `{`
	if encoder.prettify {
		encoder.increment_level()
		encoder.add_indent()
	}

	mut i := 0
	for key, value in val {
		encoder.encode_string('${key}')
		encoder.output << `:`
		if encoder.prettify {
			encoder.output << ` `
		}
		encoder.encode_value[T](value)
		if i < val.len - 1 {
			encoder.output << `,`
			if encoder.prettify {
				encoder.add_indent()
			}
		} else {
			if encoder.prettify {
				encoder.decrement_level()
				encoder.add_indent()
			}
		}

		i++
	}

	encoder.output << `}`
}

fn (mut encoder Encoder) encode_enum[T](val T) {
	if encoder.enum_as_int || enum_uses_json_as_number[T]() {
		encoder.encode_number(int(val))
	} else {
		mut enum_val := 'unknown enum value'
		$for member in T.values {
			if member.value == val {
				enum_val = member.name
				for attr in member.attrs {
					if json_attr := json_attr_value(attr) {
						enum_val = json_attr
					}
				}
			}
		}
		encoder.output << `"`
		unsafe { encoder.output.push_many(enum_val.str, enum_val.len) }
		encoder.output << `"`
	}
}

fn (mut encoder Encoder) encode_sumtype[T](val T) {
	$if T is $pointer {
		// Pointer types are handled by encode_value's $pointer branch;
		// this instantiation is generated but never called.
	} $else {
		$for variant in T.variants {
			if val is variant {
				variant_name := sumtype_variant_name(typeof(variant.typ).name)
				$if variant.typ is time.Time {
					if T.name in ['x.json2.Any', 'json2.Any', 'Any'] {
						variant_value := val
						encoder.encode_value(variant_value)
					} else {
						encoder.encode_sumtype_time_variant(val, variant_name)
					}
				} $else $if variant.typ is $struct {
					if T.name in ['x.json2.Any', 'json2.Any', 'Any'] {
						variant_value := val
						encoder.encode_value(variant_value)
					} else {
						encoder.encode_sumtype_struct_variant(val, variant_name)
					}
				} $else $if variant.typ is $map {
					encoder.encode_value(val)
				} $else {
					variant_value := val
					encoder.encode_value(variant_value)
				}
			}
		}
	}
}

fn (mut encoder Encoder) encode_object_key(is_first bool, key string) bool {
	if is_first {
		if encoder.prettify {
			encoder.increment_level()
		}
	} else {
		encoder.output << `,`
	}
	if encoder.prettify {
		encoder.add_indent()
	}
	encoder.encode_string(key)
	encoder.output << `:`
	if encoder.prettify {
		encoder.output << ` `
	}
	return false
}

fn (mut encoder Encoder) encode_sumtype_struct_variant[T](val T, variant_name string) {
	$for field in T.fields {
		$if field.is_embed {
			unsafe { encoder.encode_sumtype_struct_variant_with_embeds(val, variant_name) }
			return
		}
	}
	encoder.output << `{`
	mut is_first := unsafe { encoder.encode_struct_fields[T](val, true, [], '') }
	is_first = encoder.encode_object_key(is_first, '_type')
	encoder.encode_string(variant_name)
	if encoder.prettify && !is_first {
		encoder.decrement_level()
		encoder.add_indent()
	}
	encoder.output << `}`
}

@[markused]
fn (mut encoder Encoder) encode_sumtype_time_variant(val time.Time, variant_name string) {
	encoder.output << `{`
	mut is_first := true
	is_first = encoder.encode_object_key(is_first, '_type')
	encoder.encode_string(variant_name)
	is_first = encoder.encode_object_key(is_first, 'value')
	encoder.encode_number(val.unix())
	if encoder.prettify && !is_first {
		encoder.decrement_level()
		encoder.add_indent()
	}
	encoder.output << `}`
}

struct EncoderFieldInfo {
	key_name string

	is_skip      bool
	is_omitempty bool
	is_required  bool
	is_json_null bool
}

fn get_value_from_optional[T](val ?T) T {
	return val or { T{} }
}

fn check_not_empty[T](val T) ?bool {
	$if T.indirections != 0 {
		return val != unsafe { nil }
	} $else $if T.unaliased_typ is string {
		if val == '' {
			return false
		}
	} $else $if T.unaliased_typ is $int || T.unaliased_typ is $float {
		if val == 0 {
			return false
		}
	} $else $if T.unaliased_typ is $array || T.unaliased_typ is $map {
		return val.len != 0
	} $else $if val is ?string {
		opt := ?string(val)
		if sval := opt {
			return sval != ''
		}
		return false
	} $else $if val is ?int {
		opt := ?int(val)
		if ival := opt {
			return ival != 0
		}
		return false
	} $else $if val is ?f64 {
		opt := ?f64(val)
		if fval := opt {
			return fval != 0.0
		}
		return false
	} $else $if val is ?f32 {
		opt := ?f32(val)
		if fval := opt {
			return fval != 0.0
		}
		return false
	}
	return true
}

// TODO: fix compilation with -autofree, and remove the tag @[manualfree] here:
@[manualfree; unsafe]
fn (mut encoder Encoder) cached_field_infos[T]() []EncoderFieldInfo {
	static field_infos := &[]EncoderFieldInfo(nil)
	if field_infos == nil {
		field_infos = &[]EncoderFieldInfo{}
		$for field in T.fields {
			mut is_skip := false
			mut key_name := ''
			mut is_omitempty := false
			mut is_required := false
			mut is_json_null := false
			for attr in field.attrs {
				match attr {
					'skip' {
						is_skip = true
						break
					}
					'omitempty' {
						is_omitempty = true
					}
					'required' {
						is_required = true
					}
					'json_null' {
						is_json_null = true
					}
					else {}
				}

				if attr.starts_with('json:') {
					json_attr := json_attr_value(attr) or { continue }
					if json_attr == '-' {
						is_skip = true
						break
					}
					key_name = json_attr
				}
			}
			field_infos << EncoderFieldInfo{
				key_name:     if key_name == '' { field.name } else { key_name }
				is_skip:      is_skip
				is_omitempty: is_omitempty
				is_required:  is_required
				is_json_null: is_json_null
			}
		}
	}
	return *field_infos
}

fn (mut encoder Encoder) encode_struct_field_value[T](val T) {
	$if T is $interface {
		encoder.encode_null()
	} $else $if T.unaliased_typ is voidptr {
		encoder.encode_null()
	} $else $if T.pointee_type is $interface {
		encoder.encode_null()
	} $else $if T is $option {
		if val == none {
			unsafe { encoder.output.push_many(null_string.str, null_string.len) }
		} else {
			encoder.encode_value(get_value_from_optional(val))
		}
	} $else $if T.indirections == 1 {
		encoder.encode_value(*val)
	} $else $if T.indirections == 2 {
		encoder.encode_value(**val)
	} $else $if T.indirections == 3 {
		encoder.encode_value(***val)
	} $else {
		encoder.encode_value(val)
	}
}

fn struct_field_is_none[T](val T) bool {
	$if T is $option {
		return val == none
	}
	return false
}

fn struct_field_is_nil[T](val T) bool {
	$if T.indirections != 0 {
		return val == unsafe { nil }
	}
	return false
}

fn struct_field_should_encode[T](field_info EncoderFieldInfo, val T) bool {
	if field_info.is_skip {
		return false
	}
	if field_info.is_omitempty {
		if !(check_not_empty(val) or { false }) {
			return false
		}
	}
	if !field_info.is_required && !field_info.is_json_null && struct_field_is_none(val) {
		return false
	}
	if struct_field_is_nil(val) {
		return false
	}
	return true
}

@[unsafe]
fn (mut encoder Encoder) encode_struct_with_embeds[T](val T) {
	encoder.output << `{`
	is_first := encoder.encode_embedded_struct_fields[T](val, true, [], [], '')
	if encoder.prettify && !is_first {
		encoder.decrement_level()
		encoder.add_indent()
	}
	encoder.output << `}`
}

@[unsafe]
fn (mut encoder Encoder) encode_sumtype_struct_variant_with_embeds[T](val T, variant_name string) {
	encoder.output << `{`
	mut is_first := encoder.encode_embedded_struct_fields[T](val, true, [], [], '')
	is_first = encoder.encode_object_key(is_first, '_type')
	encoder.encode_string(variant_name)
	if encoder.prettify && !is_first {
		encoder.decrement_level()
		encoder.add_indent()
	}
	encoder.output << `}`
}

@[unsafe]
fn (mut encoder Encoder) encode_struct_fields[T](val T, was_first bool, old_used_keys []string, prefix string) bool {
	field_infos := encoder.cached_field_infos[T]()
	mut is_first := was_first
	mut used_keys := old_used_keys
	mut i := 0

	$for field in T.fields {
		$if !field.is_embed {
			field_info := field_infos[i]
			mut write_field := true

			$if field.typ is $shared {
				shared field_value := unsafe { val.$(field.name) }
				rlock field_value {
					write_field = struct_field_should_encode(field_info, field_value)

					if write_field {
						if field_info.key_name in old_used_keys {
							is_first = encoder.encode_object_key(is_first, prefix +
								field_info.key_name)
						} else {
							is_first = encoder.encode_object_key(is_first, field_info.key_name)
							used_keys << field_info.key_name
						}
						encoder.encode_struct_field_value(field_value)
					}
				}
			} $else {
				write_field = struct_field_should_encode(field_info, val.$(field.name))

				if write_field {
					if field_info.key_name in old_used_keys {
						is_first = encoder.encode_object_key(is_first, prefix + field_info.key_name)
					} else {
						is_first = encoder.encode_object_key(is_first, field_info.key_name)
						used_keys << field_info.key_name
					}
					encoder.encode_struct_field_value(val.$(field.name))
				}
			}
		}
		i++
	}
	$for field in T.fields {
		$if field.is_embed {
			new_prefix := prefix + field.name + '.'
			$if field.typ is $shared {
				shared field_value := unsafe { val.$(field.name) }
				rlock field_value {
					is_first = encoder.encode_struct_fields(field_value, is_first, used_keys,
						new_prefix)
				}
			} $else {
				is_first = encoder.encode_struct_fields(val.$(field.name), is_first, used_keys,
					new_prefix)
			}
		}
	}
	return is_first
}

@[unsafe]
fn (mut encoder Encoder) encode_embedded_struct_fields[T](val T, was_first bool, old_used_keys []string, reserved_keys []string, prefix string) bool {
	field_infos := encoder.cached_field_infos[T]()
	mut is_first := was_first
	mut used_keys := old_used_keys.clone()
	mut i := 0

	$for field in T.fields {
		$if field.is_embed {
			mut child_reserved_keys := reserved_keys.clone()
			mut reserved_i := 0
			$for reserved_field in T.fields {
				reserved_field_info := field_infos[reserved_i]
				$if !reserved_field.is_embed {
					if !reserved_field_info.is_skip {
						child_reserved_keys << reserved_field_info.key_name
					}
				}
				reserved_i++
			}
			new_prefix := prefix + field.name + '.'
			$if field.typ is $shared {
				shared field_value := unsafe { val.$(field.name) }
				rlock field_value {
					is_first = encoder.encode_embedded_struct_fields(field_value, is_first,
						used_keys, child_reserved_keys, new_prefix)
				}
			} $else {
				is_first = encoder.encode_embedded_struct_fields(val.$(field.name), is_first,
					used_keys, child_reserved_keys, new_prefix)
			}
		} $else {
			field_info := field_infos[i]
			mut write_field := true
			$if field.typ is $shared {
				shared field_value := unsafe { val.$(field.name) }
				rlock field_value {
					write_field = struct_field_should_encode(field_info, field_value)
					if write_field {
						should_prefix := field_info.key_name in used_keys
							|| field_info.key_name in reserved_keys
						json_key := if should_prefix {
							prefix + field_info.key_name
						} else {
							field_info.key_name
						}
						is_first = encoder.encode_object_key(is_first, json_key)
						if !should_prefix {
							used_keys << field_info.key_name
						}
						encoder.encode_struct_field_value(field_value)
					}
				}
			} $else {
				write_field = struct_field_should_encode(field_info, val.$(field.name))
				if write_field {
					should_prefix := field_info.key_name in used_keys
						|| field_info.key_name in reserved_keys
					json_key := if should_prefix {
						prefix + field_info.key_name
					} else {
						field_info.key_name
					}
					is_first = encoder.encode_object_key(is_first, json_key)
					if !should_prefix {
						used_keys << field_info.key_name
					}
					encoder.encode_struct_field_value(val.$(field.name))
				}
			}
		}
		i++
	}
	return is_first
}

fn (mut encoder Encoder) encode_custom[T](val T) {
	integer_val := val.to_json()
	unsafe { encoder.output.push_many(integer_val.str, integer_val.len) }
}

fn (mut encoder Encoder) encode_custom2[T](val T) {
	integer_val := val.json_str()
	unsafe { encoder.output.push_many(integer_val.str, integer_val.len) }
}

fn (mut encoder Encoder) increment_level() {
	encoder.level++
	encoder.prefix = encoder.newline_string + encoder.indent_string.repeat(encoder.level)
}

fn (mut encoder Encoder) decrement_level() {
	encoder.level--
	encoder.prefix = encoder.newline_string + encoder.indent_string.repeat(encoder.level)
}

fn (mut encoder Encoder) add_indent() {
	unsafe { encoder.output.push_many(encoder.prefix.str, encoder.prefix.len) }
}