package key_codec

import "core:bytes"
import "core:encoding/varint"
import "core:mem"

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

// Encode a varint length prefix
encode_varint :: proc(buf: ^bytes.Buffer, value: int) {
	temp: [10]byte
	n := varint.encode_u64(temp[:], u64(value))
	bytes.buffer_write(buf, temp[:n])
}

// Decode a varint length prefix
decode_varint :: proc(data: []byte, offset: ^int) -> (value: int, ok: bool) {
	if offset^ >= len(data) {
		return 0, false
	}
	
	val, n := varint.decode_u64(data[offset^:])
	if n <= 0 {
		return 0, false
	}
	
	offset^ += n
	return int(val), 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?]
build_gsi_key :: proc(table_name: string, index_name: string, gsi_pk: []byte, gsi_sk: 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.GSI))
	
	// 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 GSI partition key
	encode_varint(&buf, len(gsi_pk))
	bytes.buffer_write(&buf, gsi_pk)
	
	// Write GSI sort key if present
	if sk, ok := gsi_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_init :: proc(data: []byte) -> Key_Decoder {
	return Key_Decoder{data = data, pos = 0}
}

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 := decoder_init(key)
	
	// 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
}