package dynamodb

import "core:bytes"

// Entity type prefix bytes for namespacing
Entity_Type :: enum u8 {
	Meta = 0x01,  // Table metadata
	Data = 0x02,  // Item data
	GSI  = 0x03,  // Global secondary index
	LSI  = 0x04,  // Local secondary index
}


// Decode a varint length prefix from a byte slice.
// Reads starting at data[offset^] and advances offset^ past the varint on success.
decode_varint :: proc(data: []byte, offset: ^int) -> (value: int, ok: bool) {
	i := offset^
	if i < 0 || i >= len(data) {
		return 0, false
	}

	u: u64 = 0
	shift: u32 = 0

	for {
		if i >= len(data) {
			return 0, false // truncated
		}

		b := data[i]
		i += 1

		u |= u64(b & 0x7F) << shift

		if (b & 0x80) == 0 {
			break
		}

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

	// ensure it fits in int on this platform
	max_int := int((~uint(0)) >> 1)
	if u > u64(max_int) {
		return 0, false
	}

	offset^ = i
	return int(u), true
}





// Build metadata key: [meta][table_name]
build_meta_key :: proc(table_name: string) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 256, context.allocator)

	// Write entity type
	bytes.buffer_write_byte(&buf, u8(Entity_Type.Meta))

	// Write table name with length prefix
	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	return bytes.buffer_to_bytes(&buf)
}

// Build data key: [data][table_name][pk_value][sk_value?]
build_data_key :: proc(table_name: string, pk_value: []byte, sk_value: Maybe([]byte)) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 512, context.allocator)

	// Write entity type
	bytes.buffer_write_byte(&buf, u8(Entity_Type.Data))

	// Write table name
	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	// Write partition key
	encode_varint(&buf, len(pk_value))
	bytes.buffer_write(&buf, pk_value)

	// Write sort key if present
	if sk, ok := sk_value.?; ok {
		encode_varint(&buf, len(sk))
		bytes.buffer_write(&buf, sk)
	}

	return bytes.buffer_to_bytes(&buf)
}

// Build table prefix for scanning: [data][table_name]
build_table_prefix :: proc(table_name: string) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 256, context.allocator)

	// Write entity type
	bytes.buffer_write_byte(&buf, u8(Entity_Type.Data))

	// Write table name
	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	return bytes.buffer_to_bytes(&buf)
}

// Build partition prefix for querying: [data][table_name][pk_value]
build_partition_prefix :: proc(table_name: string, pk_value: []byte) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 512, context.allocator)

	// Write entity type
	bytes.buffer_write_byte(&buf, u8(Entity_Type.Data))

	// Write table name
	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	// Write partition key
	encode_varint(&buf, len(pk_value))
	bytes.buffer_write(&buf, pk_value)

	return bytes.buffer_to_bytes(&buf)
}

// Build GSI key: [gsi][table_name][index_name][gsi_pk][gsi_sk?][base_pk][base_sk?]
build_gsi_key :: proc(
	table_name: string,
	index_name: string,
	gsi_pk: []byte,
	gsi_sk: Maybe([]byte),
	base_pk: []byte,
	base_sk: Maybe([]byte),
) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 512, context.allocator)

	bytes.buffer_write_byte(&buf, u8(Entity_Type.GSI))

	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	encode_varint(&buf, len(index_name))
	bytes.buffer_write_string(&buf, index_name)

	encode_varint(&buf, len(gsi_pk))
	bytes.buffer_write(&buf, gsi_pk)

	if sk, ok := gsi_sk.?; ok {
		encode_varint(&buf, len(sk))
		bytes.buffer_write(&buf, sk)
	}

	// tie-breaker: base table primary key
	encode_varint(&buf, len(base_pk))
	bytes.buffer_write(&buf, base_pk)

	if sk, ok := base_sk.?; ok {
		encode_varint(&buf, len(sk))
		bytes.buffer_write(&buf, sk)
	}

	return bytes.buffer_to_bytes(&buf)
}

// Build LSI key: [lsi][table_name][index_name][pk][lsi_sk]
build_lsi_key :: proc(table_name: string, index_name: string, pk: []byte, lsi_sk: []byte) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 512, context.allocator)

	// Write entity type
	bytes.buffer_write_byte(&buf, u8(Entity_Type.LSI))

	// Write table name
	encode_varint(&buf, len(table_name))
	bytes.buffer_write_string(&buf, table_name)

	// Write index name
	encode_varint(&buf, len(index_name))
	bytes.buffer_write_string(&buf, index_name)

	// Write partition key
	encode_varint(&buf, len(pk))
	bytes.buffer_write(&buf, pk)

	// Write LSI sort key
	encode_varint(&buf, len(lsi_sk))
	bytes.buffer_write(&buf, lsi_sk)

	return bytes.buffer_to_bytes(&buf)
}

// Key decoder for reading binary keys
Key_Decoder :: struct {
	data: []byte,
	pos:  int,
}

decoder_read_entity_type :: proc(decoder: ^Key_Decoder) -> (Entity_Type, bool) {
	if decoder.pos >= len(decoder.data) {
		return .Meta, false
	}

	entity_type := Entity_Type(decoder.data[decoder.pos])
	decoder.pos += 1
	return entity_type, true
}

decoder_read_segment :: proc(decoder: ^Key_Decoder) -> (segment: []byte, ok: bool) {
	// Read length
	length := decode_varint(decoder.data, &decoder.pos) or_return

	// Read data
	if decoder.pos + length > len(decoder.data) {
		return nil, false
	}

	// Return slice (owned by caller via context.allocator)
	segment = make([]byte, length, context.allocator)
	copy(segment, decoder.data[decoder.pos:decoder.pos + length])
	decoder.pos += length

	return segment, true
}

decoder_read_segment_borrowed :: proc(decoder: ^Key_Decoder) -> (segment: []byte, ok: bool) {
	// Read length
	length := decode_varint(decoder.data, &decoder.pos) or_return

	// Return borrowed slice
	if decoder.pos + length > len(decoder.data) {
		return nil, false
	}

	segment = decoder.data[decoder.pos:decoder.pos + length]
	decoder.pos += length

	return segment, true
}

decoder_has_more :: proc(decoder: ^Key_Decoder) -> bool {
	return decoder.pos < len(decoder.data)
}

// Decode a data key back into components
Decoded_Data_Key :: struct {
	table_name: string,
	pk_value:   []byte,
	sk_value:   Maybe([]byte),
}

decode_data_key :: proc(key: []byte) -> (result: Decoded_Data_Key, ok: bool) {
	decoder := Key_Decoder{data = key, pos = 0}

	// Read and verify entity type
	entity_type := decoder_read_entity_type(&decoder) or_return
	if entity_type != .Data {
		return {}, false
	}

	// Read table name
	table_name_bytes := decoder_read_segment(&decoder) or_return
	result.table_name = string(table_name_bytes)

	// Read partition key
	result.pk_value = decoder_read_segment(&decoder) or_return

	// Read sort key if present
	if decoder_has_more(&decoder) {
		sk := decoder_read_segment(&decoder) or_return
		result.sk_value = sk
	}

	return result, true
}