// Binary TLV (Type-Length-Value) encoding for DynamoDB items
// Replaces JSON storage with efficient binary format
// Format: [attribute_count][name_len][name][type_tag][value_len][value]...
package dynamodb

import "core:bytes"
import "core:slice"

// Type tags for binary encoding (1 byte each)
Type_Tag :: enum u8 {
	// Scalar types
	String  = 0x01, // S
	Number  = 0x02, // N (stored as string)
	Binary  = 0x03, // B (base64 string)
	Boolean = 0x04, // BOOL
	Null    = 0x05, // NULL

	// Set types
	String_Set = 0x10, // SS
	Number_Set = 0x11, // NS
	Binary_Set = 0x12, // BS

	// Complex types
	List = 0x20, // L
	Map  = 0x21, // M
}

// ============================================================================
// Encoding (Item → Binary)
// ============================================================================

// Encode an Item to binary TLV format
// Format: [attribute_count:varint][attributes...]
// Each attribute: [name_len:varint][name:bytes][type_tag:u8][value_encoded:bytes]
encode :: proc(item: Item) -> ([]byte, bool) {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 1024, context.allocator)
	defer bytes.buffer_destroy(&buf)

	// Write attribute count
	encode_varint(&buf, len(item))

	// Collect and sort keys for deterministic encoding
	keys := make([dynamic]string, context.temp_allocator)
	for key in item {
		append(&keys, key)
	}

	slice.sort_by(keys[:], proc(a, b: string) -> bool {
		return a < b
	})

	// Encode each attribute
	for key in keys {
		value := item[key]

		// Write attribute name
		encode_varint(&buf, len(key))
		bytes.buffer_write_string(&buf, key)

		// Encode attribute value
		ok := encode_attribute_value(&buf, value)
		if !ok {
			return nil, false
		}
	}

	return bytes.buffer_to_bytes(&buf), true
}

// Encode an AttributeValue to binary format
encode_attribute_value :: proc(buf: ^bytes.Buffer, attr: Attribute_Value) -> bool {
	switch v in attr {
	case String:
		bytes.buffer_write_byte(buf, u8(Type_Tag.String))
		encode_varint(buf, len(v))
		bytes.buffer_write_string(buf, string(v))

	case DDB_Number:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Number))
		// Store as string in item encoding
		num_str := format_ddb_number(v)
		encode_varint(buf, len(num_str))
		bytes.buffer_write_string(buf, num_str)

	case Number:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Number))
		encode_varint(buf, len(v))
		bytes.buffer_write_string(buf, string(v))

	case Binary:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Binary))
		encode_varint(buf, len(v))
		bytes.buffer_write_string(buf, string(v))

	case Bool:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Boolean))
		bytes.buffer_write_byte(buf, 1 if bool(v) else 0)

	case Null:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Null))
		// NULL has no value bytes

	case DDB_Number_Set:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Number_Set))  // Use Number_Set tag, not DDB_Number_Set
		encode_varint(buf, len(v))
		for num in v {
			// Format the DDB_Number to a string
			num_str := format_ddb_number(num)
			encode_varint(buf, len(num_str))
			bytes.buffer_write_string(buf, num_str)
		}

	case String_Set:
		bytes.buffer_write_byte(buf, u8(Type_Tag.String_Set))
		encode_varint(buf, len(v))
		for s in v {
			encode_varint(buf, len(s))
			bytes.buffer_write_string(buf, s)
		}

	case Number_Set:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Number_Set))
		encode_varint(buf, len(v))
		for n in v {
			encode_varint(buf, len(n))
			bytes.buffer_write_string(buf, n)
		}

	case Binary_Set:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Binary_Set))
		encode_varint(buf, len(v))
		for b in v {
			encode_varint(buf, len(b))
			bytes.buffer_write_string(buf, b)
		}

	case List:
		bytes.buffer_write_byte(buf, u8(Type_Tag.List))
		encode_varint(buf, len(v))
		for item in v {
			ok := encode_attribute_value(buf, item)
			if !ok {
				return false
			}
		}

	case Map:
		bytes.buffer_write_byte(buf, u8(Type_Tag.Map))
		encode_varint(buf, len(v))

		// Collect and sort keys for deterministic encoding
		keys := make([dynamic]string, context.temp_allocator)
		for key in v {
			append(&keys, key)
		}

		slice.sort_by(keys[:], proc(a, b: string) -> bool {
			return a < b
		})

		// Encode each map entry
		for key in keys {
			value := v[key]
			encode_varint(buf, len(key))
			bytes.buffer_write_string(buf, key)
			ok := encode_attribute_value(buf, value)
			if !ok {
				return false
			}
		}
	}

	return true
}

// ============================================================================
// Decoding (Binary → Item)
// ============================================================================

// Binary decoder helper
Binary_Decoder :: struct {
	data: []byte,
	pos:  int,
}

decoder_read_byte :: proc(decoder: ^Binary_Decoder) -> (u8, bool) {
	if decoder.pos >= len(decoder.data) {
		return 0, false
	}

	byte := decoder.data[decoder.pos]
	decoder.pos += 1
	return byte, true
}

decoder_read_bytes :: proc(decoder: ^Binary_Decoder, length: int) -> ([]byte, bool) {
	if decoder.pos + length > len(decoder.data) {
		return nil, false
	}

	bytes := decoder.data[decoder.pos:decoder.pos + length]
	decoder.pos += length
	return bytes, true
}

decoder_read_varint :: proc(decoder: ^Binary_Decoder) -> (int, bool) {
	result: int = 0
	shift: uint = 0

	for decoder.pos < len(decoder.data) {
		byte := decoder.data[decoder.pos]
		decoder.pos += 1

		result |= int(byte & 0x7F) << shift

		if (byte & 0x80) == 0 {
			return result, true
		}

		shift += 7
		if shift >= 64 {
			return 0, false // Varint overflow
		}
	}

	return 0, false // Unexpected end of data
}

// Decode binary TLV format back into an Item
decode :: proc(data: []byte) -> (Item, bool) {
	decoder :=  Binary_Decoder{data = data, pos = 0}

	attr_count, count_ok := decoder_read_varint(&decoder)
	if !count_ok {
		return {}, false
	}

	item := make(Item)

	for _ in 0..<attr_count {
		// Read attribute name
		name_len, name_len_ok := decoder_read_varint(&decoder)
		if !name_len_ok {
			// Cleanup on error
			item_destroy(&item)
			return {}, false
		}

		name_bytes, name_ok := decoder_read_bytes(&decoder, name_len)
		if !name_ok {
			item_destroy(&item)
			return {}, false
		}

		owned_name := string(name_bytes)
		owned_name = transmute(string)slice.clone(transmute([]byte)owned_name)

		// Read attribute value
		value, value_ok := decode_attribute_value(&decoder)
		if !value_ok {
			delete(owned_name)
			item_destroy(&item)
			return {}, false
		}

		item[owned_name] = value
	}

	return item, true
}

// Decode an AttributeValue from binary format
decode_attribute_value :: proc(decoder: ^Binary_Decoder) -> (Attribute_Value, bool) {
	type_byte, type_ok := decoder_read_byte(decoder)
	if !type_ok {
		return nil, false
	}

	type_tag := Type_Tag(type_byte)

	switch type_tag {
	case .String:
		length, len_ok := decoder_read_varint(decoder)
		if !len_ok {
			return nil, false
		}

		data, data_ok := decoder_read_bytes(decoder, length)
		if !data_ok {
			return nil, false
		}

		str := string(data)
		owned := transmute(string)slice.clone(transmute([]byte)str)
		return String(owned), true

	case .Number:
		length, len_ok := decoder_read_varint(decoder)
		if !len_ok {
			return nil, false
		}

		data, data_ok := decoder_read_bytes(decoder, length)
		if !data_ok {
			return nil, false
		}

		num_str := string(data)

		// Parse into DDB_Number
		ddb_num, num_ok := parse_ddb_number(num_str)
		if !num_ok {
			return nil, false
		}

		return ddb_num, true

	case .Binary:
		length, len_ok := decoder_read_varint(decoder)
		if !len_ok {
			return nil, false
		}

		data, data_ok := decoder_read_bytes(decoder, length)
		if !data_ok {
			return nil, false
		}

		str := string(data)
		owned := transmute(string)slice.clone(transmute([]byte)str)
		return Binary(owned), true

	case .Boolean:
		byte, byte_ok := decoder_read_byte(decoder)
		if !byte_ok {
			return nil, false
		}

		return Bool(byte != 0), true

	case .Null:
		return Null(true), true

	case .String_Set:
		count, count_ok := decoder_read_varint(decoder)
		if !count_ok {
			return nil, false
		}

		strings := make([]string, count)

		for i in 0..<count {
			length, len_ok := decoder_read_varint(decoder)
			if !len_ok {
				// Cleanup on error
				for j in 0..<i {
					delete(strings[j])
				}
				delete(strings)
				return nil, false
			}

			data, data_ok := decoder_read_bytes(decoder, length)
			if !data_ok {
				for j in 0..<i {
					delete(strings[j])
				}
				delete(strings)
				return nil, false
			}

			str := string(data)
			strings[i] = transmute(string)slice.clone(transmute([]byte)str)
		}

		return String_Set(strings), true

	case .Number_Set:
		count, count_ok := decoder_read_varint(decoder)
		if !count_ok {
			return nil, false
		}

		numbers := make([]DDB_Number, count)  // Changed to DDB_Number

		for i in 0..<count {
			length, len_ok := decoder_read_varint(decoder)
			if !len_ok {
				// No cleanup needed for DDB_Number (no heap allocations)
				delete(numbers)
				return nil, false
			}

			data, data_ok := decoder_read_bytes(decoder, length)
			if !data_ok {
				delete(numbers)
				return nil, false
			}

			num_str := string(data)

			// Parse into DDB_Number
			ddb_num, num_ok := parse_ddb_number(num_str)
			if !num_ok {
				delete(numbers)
				return nil, false
			}

			numbers[i] = ddb_num
		}

		return DDB_Number_Set(numbers), true

	case .Binary_Set:
		count, count_ok := decoder_read_varint(decoder)
		if !count_ok {
			return nil, false
		}

		binaries := make([]string, count)

		for i in 0..<count {
			length, len_ok := decoder_read_varint(decoder)
			if !len_ok {
				for j in 0..<i {
					delete(binaries[j])
				}
				delete(binaries)
				return nil, false
			}

			data, data_ok := decoder_read_bytes(decoder, length)
			if !data_ok {
				for j in 0..<i {
					delete(binaries[j])
				}
				delete(binaries)
				return nil, false
			}

			str := string(data)
			binaries[i] = transmute(string)slice.clone(transmute([]byte)str)
		}

		return Binary_Set(binaries), true

	case .List:
		count, count_ok := decoder_read_varint(decoder)
		if !count_ok {
			return nil, false
		}

		list := make([]Attribute_Value, count)

		for i in 0..<count {
			value, value_ok := decode_attribute_value(decoder)
			if !value_ok {
				// Cleanup on error
				for j in 0..<i {
					item := list[j]
					attr_value_destroy(&item)
				}
				delete(list)
				return nil, false
			}

			list[i] = value
		}

		return List(list), true

	case .Map:
		count, count_ok := decoder_read_varint(decoder)
		if !count_ok {
			return nil, false
		}

		attr_map := make(map[string]Attribute_Value)

		for _ in 0..<count {
			// Read key
			key_len, key_len_ok := decoder_read_varint(decoder)
			if !key_len_ok {
				// Cleanup on error
				for k, v in attr_map {
					delete(k)
					v_copy := v
					attr_value_destroy(&v_copy)
				}
				delete(attr_map)
				return nil, false
			}

			key_bytes, key_ok := decoder_read_bytes(decoder, key_len)
			if !key_ok {
				for k, v in attr_map {
					delete(k)
					v_copy := v
					attr_value_destroy(&v_copy)
				}
				delete(attr_map)
				return nil, false
			}

			key := string(key_bytes)
			owned_key := transmute(string)slice.clone(transmute([]byte)key)

			// Read value
			value, value_ok := decode_attribute_value(decoder)
			if !value_ok {
				delete(owned_key)
				for k, v in attr_map {
					delete(k)
					v_copy := v
					attr_value_destroy(&v_copy)
				}
				delete(attr_map)
				return nil, false
			}

			attr_map[owned_key] = value
		}

		return Map(attr_map), true
	}

	return nil, false
}

// ============================================================================
// Varint Encoding (Encodes a varint length prefix)
// ============================================================================

encode_varint :: proc(buf: ^bytes.Buffer, value: int) {
	v := value
	for {
		byte := u8(v & 0x7F)
		v >>= 7

		if v == 0 {
			bytes.buffer_write_byte(buf, byte)
			return
		} else {
			bytes.buffer_write_byte(buf, byte | 0x80)
		}
	}
}