package dynamodb

import "core:fmt"
import "core:strconv"
import "core:strings"
import "core:bytes"

// ============================================================================
// DynamoDB Number Type
//
// DynamoDB numbers are arbitrary-precision decimals with up to 38 digits of
// precision. They can be positive, negative, or zero.
//
// We store numbers internally as:
// - sign: bool (true = positive/zero, false = negative)
// - integer_part: string (digits only, no sign)
// - fractional_part: string (digits only, if any)
// - exponent: i32 (for scientific notation, if needed)
//
// This preserves the original precision and allows proper ordering.
// ============================================================================

DDB_Number :: struct {
	sign:            bool,   // true = positive/zero, false = negative
	integer_part:    string, // digits only (e.g., "123")
	fractional_part: string, // digits only (e.g., "456" for .456)
	exponent:        i32,    // scientific notation exponent (usually 0)
}

// Parse a number string into DDB_Number
// Supports formats: "123", "-123", "123.456", "1.23e10", "-1.23e-5"
parse_ddb_number :: proc(s: string) -> (DDB_Number, bool) {
	if len(s) == 0 {
		return {}, false
	}

	num: DDB_Number
	str := s

	// Parse sign
	if str[0] == '-' {
		num.sign = false
		str = str[1:]
	} else if str[0] == '+' {
		num.sign = true
		str = str[1:]
	} else {
		num.sign = true
	}

	if len(str) == 0 {
		return {}, false
	}

	// Find exponent if present (e or E)
	exp_pos := -1
	for i in 0..<len(str) {
		if str[i] == 'e' || str[i] == 'E' {
			exp_pos = i
			break
		}
	}

	// Parse mantissa
	mantissa := str
	if exp_pos >= 0 {
		mantissa = str[:exp_pos]
		exp_str := str[exp_pos+1:]
		exp_val, exp_ok := strconv.parse_i64(exp_str)
		if !exp_ok {
			return {}, false
		}
		num.exponent = i32(exp_val)
	}

	// Find decimal point
	dot_pos := -1
	for i in 0..<len(mantissa) {
		if mantissa[i] == '.' {
			dot_pos = i
			break
		}
	}

	// Parse integer and fractional parts
	if dot_pos >= 0 {
		num.integer_part = mantissa[:dot_pos]
		num.fractional_part = mantissa[dot_pos+1:]

		// Validate fractional part
		for c in num.fractional_part {
			if c < '0' || c > '9' {
				return {}, false
			}
		}
	} else {
		num.integer_part = mantissa
	}

	// Validate integer part (at least one digit, all digits)
	if len(num.integer_part) == 0 {
		num.integer_part = "0"
	}
	for c in num.integer_part {
		if c < '0' || c > '9' {
			return {}, false
		}
	}

	// Normalize: remove leading zeros from integer part (except if it's just "0")
	num = normalize_ddb_number(num)

	// Check precision (DynamoDB supports up to 38 digits)
	total_digits := len(num.integer_part) + len(num.fractional_part)
	if total_digits > 38 {
		return {}, false
	}

	// Special case: if the number is zero
	if is_ddb_number_zero(num) {
		num.sign = true
		num.exponent = 0
	}

	return num, true
}

// Normalize a DDB_Number (remove leading zeros, trailing fractional zeros)
normalize_ddb_number :: proc(num: DDB_Number) -> DDB_Number {
	result := num

	// Remove leading zeros from integer part
	int_part := num.integer_part
	for len(int_part) > 1 && int_part[0] == '0' {
		int_part = int_part[1:]
	}
	result.integer_part = int_part

	// Remove trailing zeros from fractional part
	frac_part := num.fractional_part
	for len(frac_part) > 0 && frac_part[len(frac_part)-1] == '0' {
		frac_part = frac_part[:len(frac_part)-1]
	}
	result.fractional_part = frac_part

	return result
}

// Check if a DDB_Number represents zero
is_ddb_number_zero :: proc(num: DDB_Number) -> bool {
	// Check if integer part is all zeros
	for c in num.integer_part {
		if c != '0' {
			return false
		}
	}
	// Check if fractional part is all zeros
	for c in num.fractional_part {
		if c != '0' {
			return false
		}
	}
	return true
}

// Convert DDB_Number to string representation
ddb_number_to_string :: proc(num: DDB_Number) -> string {
	builder := strings.builder_make()

	if !num.sign {
		strings.write_string(&builder, "-")
	}

	strings.write_string(&builder, num.integer_part)

	if len(num.fractional_part) > 0 {
		strings.write_string(&builder, ".")
		strings.write_string(&builder, num.fractional_part)
	}

	if num.exponent != 0 {
		fmt.sbprintf(&builder, "e%d", num.exponent)
	}

	return strings.to_string(builder)
}

// Compare two DDB_Numbers
// Returns: -1 if a < b, 0 if a == b, 1 if a > b
compare_ddb_numbers :: proc(a: DDB_Number, b: DDB_Number) -> int {
	// Handle zero cases
	a_zero := is_ddb_number_zero(a)
	b_zero := is_ddb_number_zero(b)

	if a_zero && b_zero {
		return 0
	}
	if a_zero {
		return b.sign ? -1 : 1  // 0 < positive, 0 > negative
	}
	if b_zero {
		return a.sign ? 1 : -1  // positive > 0, negative < 0
	}

	// Different signs
	if a.sign != b.sign {
		return a.sign ? 1 : -1  // positive > negative
	}

	// Same sign - compare magnitudes
	mag_cmp := compare_ddb_number_magnitudes(a, b)

	// If negative, reverse the comparison
	if !a.sign {
		return -mag_cmp
	}
	return mag_cmp
}

// Compare magnitudes (absolute values) of two DDB_Numbers
compare_ddb_number_magnitudes :: proc(a: DDB_Number, b: DDB_Number) -> int {
	// Adjust for exponents first
	a_adj := adjust_for_exponent(a)
	b_adj := adjust_for_exponent(b)

	// Compare integer parts length
	if len(a_adj.integer_part) != len(b_adj.integer_part) {
		return len(a_adj.integer_part) > len(b_adj.integer_part) ? 1 : -1
	}

	// Compare integer parts digit by digit
	for i in 0..<len(a_adj.integer_part) {
		if a_adj.integer_part[i] != b_adj.integer_part[i] {
			return a_adj.integer_part[i] > b_adj.integer_part[i] ? 1 : -1
		}
	}

	// Integer parts equal, compare fractional parts
	max_frac_len := max(len(a_adj.fractional_part), len(b_adj.fractional_part))
	for i in 0..<max_frac_len {
		a_digit := i < len(a_adj.fractional_part) ? a_adj.fractional_part[i] : '0'
		b_digit := i < len(b_adj.fractional_part) ? b_adj.fractional_part[i] : '0'

		if a_digit != b_digit {
			return a_digit > b_digit ? 1 : -1
		}
	}

	return 0
}

// Adjust a number for its exponent (conceptually multiply by 10^exponent)
adjust_for_exponent :: proc(num: DDB_Number) -> DDB_Number {
	if num.exponent == 0 {
		return num
	}

	result := num
	result.exponent = 0

	if num.exponent > 0 {
		// Shift decimal point right
		exp := int(num.exponent)
		frac := num.fractional_part

		// Move fractional digits to integer part
		shift := min(exp, len(frac))
		result.integer_part = strings.concatenate({num.integer_part, frac[:shift]})
		result.fractional_part = frac[shift:]

		// Add zeros if needed
		if exp > len(frac) {
			zeros := strings.repeat("0", exp - len(frac))
			result.integer_part = strings.concatenate({result.integer_part, zeros})
		}
	} else {
		// Shift decimal point left
		exp := -int(num.exponent)
		int_part := num.integer_part

		// Move integer digits to fractional part
		shift := min(exp, len(int_part))
		result.integer_part = int_part[:len(int_part)-shift]
		if len(result.integer_part) == 0 {
			result.integer_part = "0"
		}
		result.fractional_part = strings.concatenate({
			int_part[len(int_part)-shift:],
			num.fractional_part,
		})

		// Add leading zeros if needed
		if exp > len(int_part) {
			zeros := strings.repeat("0", exp - len(int_part))
			result.fractional_part = strings.concatenate({zeros, result.fractional_part})
		}
	}

	return normalize_ddb_number(result)
}

// ============================================================================
// Canonical Encoding for Sort Keys
//
// For numbers to sort correctly in byte-wise comparisons, we need a
// canonical encoding that preserves numeric ordering.
//
// Encoding format:
// - 1 byte: sign/magnitude marker
//   - 0x00: negative infinity (reserved)
//   - 0x01-0x7F: negative numbers (inverted magnitude)
//   - 0x80: zero
//   - 0x81-0xFE: positive numbers (magnitude)
//   - 0xFF: positive infinity (reserved)
// - N bytes: encoded magnitude (variable length)
//
// For positive numbers: we encode the magnitude directly with leading byte
// indicating number of integer digits.
//
// For negative numbers: we encode the magnitude inverted (bitwise NOT) so
// that larger negative numbers sort before smaller ones.
// ============================================================================

// Encode a DDB_Number into canonical byte form for sort keys
encode_ddb_number_for_sort :: proc(num: DDB_Number) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 64, context.allocator)

	if is_ddb_number_zero(num) {
		bytes.buffer_write_byte(&buf, 0x80)
		return bytes.buffer_to_bytes(&buf)
	}

	// Get normalized magnitude
	norm := normalize_ddb_number(num)
	adj := adjust_for_exponent(norm)

	// Encode magnitude bytes
	mag_bytes := encode_magnitude(adj)

	if num.sign {
		// Positive number: 0x81 + magnitude
		bytes.buffer_write_byte(&buf, 0x81)
		bytes.buffer_write(&buf, mag_bytes)
	} else {
		// Negative number: 0x7F - inverted magnitude
		bytes.buffer_write_byte(&buf, 0x7F)
		// Invert all magnitude bytes
		for b in mag_bytes {
			bytes.buffer_write_byte(&buf, ~b)
		}
	}

	return bytes.buffer_to_bytes(&buf)
}

// Encode the magnitude of a number (integer + fractional parts)
encode_magnitude :: proc(num: DDB_Number) -> []byte {
	buf: bytes.Buffer
	bytes.buffer_init_allocator(&buf, 0, 32, context.allocator)

	// Write length of integer part as varint
	int_len := u64(len(num.integer_part))
	encode_varint(&buf, int_len)

	// Write integer digits
	bytes.buffer_write_string(&buf, num.integer_part)

	// Write fractional digits if any
	if len(num.fractional_part) > 0 {
		bytes.buffer_write_string(&buf, num.fractional_part)
	}

	return bytes.buffer_to_bytes(&buf)
}

// Decode a canonically encoded number back to DDB_Number
decode_ddb_number_from_sort :: proc(data: []byte) -> (DDB_Number, bool) {
	if len(data) == 0 {
		return {}, false
	}

	marker := data[0]

	// Zero
	if marker == 0x80 {
		return DDB_Number{
			sign = true,
			integer_part = "0",
			fractional_part = "",
			exponent = 0,
		}, true
	}

	// Positive number
	if marker == 0x81 {
		return decode_magnitude(data[1:], true)
	}

	// Negative number (inverted bytes)
	if marker == 0x7F {
		// Un-invert the bytes
		inverted := make([]byte, len(data)-1)
		defer delete(inverted)
		for i in 0..<len(inverted) {
			inverted[i] = ~data[i+1]
		}
		return decode_magnitude(inverted, false)
	}

	return {}, false
}

// Decode magnitude bytes back to a DDB_Number
decode_magnitude :: proc(data: []byte, positive: bool) -> (DDB_Number, bool) {
	if len(data) == 0 {
		return {}, false
	}

	// Read integer length
	int_len, bytes_read := decode_varint(data)
	if bytes_read == 0 || int_len == 0 {
		return {}, false
	}

	offset := bytes_read

	// Read integer part
	if offset + int(int_len) > len(data) {
		return {}, false
	}
	int_part := string(data[offset:offset + int(int_len)])
	offset += int(int_len)

	// Read fractional part if any
	frac_part := ""
	if offset < len(data) {
		frac_part = string(data[offset:])
	}

	return DDB_Number{
		sign = positive,
		integer_part = int_part,
		fractional_part = frac_part,
		exponent = 0,
	}, true
}

// ============================================================================
// Decimal Arithmetic (38-digit precision, no float conversion)
// ============================================================================

MAX_DDB_PRECISION :: 38

// Add two DDB_Numbers with full decimal precision.
// Returns an owned DDB_Number.
add_ddb_numbers :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	if is_ddb_number_zero(a) { return clone_ddb_number(b), true }
	if is_ddb_number_zero(b) { return clone_ddb_number(a), true }

	if a.sign == b.sign {
		// Same sign: add magnitudes, keep sign
		result, ok := add_magnitudes(a, b)
		if !ok { return {}, false }
		result.sign = a.sign
		return result, true
	}

	// Different signs: subtract smaller magnitude from larger
	cmp := compare_ddb_number_magnitudes(a, b)
	if cmp == 0 {
		return DDB_Number{
			sign            = true,
			integer_part    = strings.clone("0"),
			fractional_part = strings.clone(""),
			exponent        = 0,
		}, true
	}

	if cmp > 0 {
		result, ok := subtract_magnitudes(a, b)
		if !ok { return {}, false }
		result.sign = a.sign
		return result, true
	} else {
		result, ok := subtract_magnitudes(b, a)
		if !ok { return {}, false }
		result.sign = b.sign
		return result, true
	}
}

// Subtract two DDB_Numbers: a - b
subtract_ddb_numbers :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	neg_b := b
	neg_b.sign = !b.sign
	return add_ddb_numbers(a, neg_b)
}

// ============================================================================
// Internal arithmetic helpers
// ============================================================================

// Expand a DDB_Number to effective integer and fractional digit bytes
// with the exponent fully applied. Returns heap-allocated slices (caller frees).
@(private="file")
expand_digits :: proc(num: DDB_Number) -> (int_digits: []u8, frac_digits: []u8) {
	dp := len(num.integer_part) + int(num.exponent)
	all_len := len(num.integer_part) + len(num.fractional_part)

	if dp <= 0 {
		// Everything is fractional, need leading zeros
		frac := make([]u8, -dp + all_len)
		for i in 0..<(-dp) {
			frac[i] = '0'
		}
		for i in 0..<len(num.integer_part) {
			frac[-dp + i] = num.integer_part[i]
		}
		for i in 0..<len(num.fractional_part) {
			frac[-dp + len(num.integer_part) + i] = num.fractional_part[i]
		}

		int_d := make([]u8, 1)
		int_d[0] = '0'
		return int_d, frac
	}

	if dp >= all_len {
		// Everything is integer, may need trailing zeros
		int_d := make([]u8, dp)
		for i in 0..<len(num.integer_part) {
			int_d[i] = num.integer_part[i]
		}
		for i in 0..<len(num.fractional_part) {
			int_d[len(num.integer_part) + i] = num.fractional_part[i]
		}
		for i in all_len..<dp {
			int_d[i] = '0'
		}
		return int_d, nil
	}

	// Decimal point falls within the original integer_part
	if dp <= len(num.integer_part) {
		int_d := make([]u8, dp)
		for i in 0..<dp {
			int_d[i] = num.integer_part[i]
		}

		frac_len := (len(num.integer_part) - dp) + len(num.fractional_part)
		frac := make([]u8, frac_len)
		for i in dp..<len(num.integer_part) {
			frac[i - dp] = num.integer_part[i]
		}
		offset := len(num.integer_part) - dp
		for i in 0..<len(num.fractional_part) {
			frac[offset + i] = num.fractional_part[i]
		}
		return int_d, frac
	}

	// Decimal point falls within the original fractional_part
	frac_split := dp - len(num.integer_part)

	int_d := make([]u8, dp)
	for i in 0..<len(num.integer_part) {
		int_d[i] = num.integer_part[i]
	}
	for i in 0..<frac_split {
		int_d[len(num.integer_part) + i] = num.fractional_part[i]
	}

	remaining := len(num.fractional_part) - frac_split
	frac: []u8 = nil
	if remaining > 0 {
		frac = make([]u8, remaining)
		for i in frac_split..<len(num.fractional_part) {
			frac[i - frac_split] = num.fractional_part[i]
		}
	}
	return int_d, frac
}

// Normalize a DDB_Number that owns its strings.
// Clones the trimmed result, frees the originals.
@(private="file")
normalize_owned :: proc(num: DDB_Number) -> DDB_Number {
	norm := normalize_ddb_number(num)

	// Clone the normalized subslices BEFORE freeing originals
	new_int := strings.clone(norm.integer_part)
	new_frac := strings.clone(norm.fractional_part)

	// Free the originals
	delete(num.integer_part)
	delete(num.fractional_part)

	return DDB_Number{
		sign            = norm.sign,
		integer_part    = new_int,
		fractional_part = new_frac,
		exponent        = norm.exponent,
	}
}

// Add absolute values. Returns owned DDB_Number (sign=true).
@(private="file")
add_magnitudes :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	a_int, a_frac := expand_digits(a)
	b_int, b_frac := expand_digits(b)
	defer { delete(a_int); delete(a_frac); delete(b_int); delete(b_frac) }

	max_int  := max(len(a_int), len(b_int))
	max_frac := max(len(a_frac), len(b_frac))
	total    := max_int + max_frac

	// Build zero-padded aligned arrays
	a_aligned := make([]u8, total)
	b_aligned := make([]u8, total)
	defer { delete(a_aligned); delete(b_aligned) }

	for i in 0..<total { a_aligned[i] = '0'; b_aligned[i] = '0' }

	// Integer digits: right-aligned in [0..max_int)
	a_off := max_int - len(a_int)
	b_off := max_int - len(b_int)
	for i in 0..<len(a_int) { a_aligned[a_off + i] = a_int[i] }
	for i in 0..<len(b_int) { b_aligned[b_off + i] = b_int[i] }

	// Fractional digits: left-aligned in [max_int..total)
	for i in 0..<len(a_frac) { a_aligned[max_int + i] = a_frac[i] }
	for i in 0..<len(b_frac) { b_aligned[max_int + i] = b_frac[i] }

	// Add right-to-left
	result := make([]u8, total + 1) // +1 for carry
	carry: u8 = 0
	for i := total - 1; i >= 0; i -= 1 {
		sum := (a_aligned[i] - '0') + (b_aligned[i] - '0') + carry
		result[i + 1] = (sum % 10) + '0'
		carry = sum / 10
	}
	result[0] = carry + '0'

	// Split: decimal point is at max_int + 1 (carry slot shifts everything)
	int_end := max_int + 1
	int_str  := strings.clone(string(result[:int_end]))
	frac_str := strings.clone(string(result[int_end:]))
	delete(result)

	num := normalize_owned(DDB_Number{
		sign            = true,
		integer_part    = int_str,
		fractional_part = frac_str,
		exponent        = 0,
	})

	if len(num.integer_part) + len(num.fractional_part) > MAX_DDB_PRECISION {
		delete(num.integer_part)
		delete(num.fractional_part)
		return {}, false
	}

	return num, true
}

// Subtract absolute values: |a| - |b|, where |a| >= |b|.
// Returns owned DDB_Number (sign=true).
@(private="file")
subtract_magnitudes :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	a_int, a_frac := expand_digits(a)
	b_int, b_frac := expand_digits(b)
	defer { delete(a_int); delete(a_frac); delete(b_int); delete(b_frac) }

	max_int  := max(len(a_int), len(b_int))
	max_frac := max(len(a_frac), len(b_frac))
	total    := max_int + max_frac

	a_aligned := make([]u8, total)
	b_aligned := make([]u8, total)
	defer { delete(a_aligned); delete(b_aligned) }

	for i in 0..<total { a_aligned[i] = '0'; b_aligned[i] = '0' }

	a_off := max_int - len(a_int)
	b_off := max_int - len(b_int)
	for i in 0..<len(a_int) { a_aligned[a_off + i] = a_int[i] }
	for i in 0..<len(b_int) { b_aligned[b_off + i] = b_int[i] }
	for i in 0..<len(a_frac) { a_aligned[max_int + i] = a_frac[i] }
	for i in 0..<len(b_frac) { b_aligned[max_int + i] = b_frac[i] }

	// Subtract right-to-left
	result := make([]u8, total)
	borrow: u8 = 0
	for i := total - 1; i >= 0; i -= 1 {
		ad := a_aligned[i] - '0'
		bd := (b_aligned[i] - '0') + borrow
		if ad < bd {
			ad += 10
			borrow = 1
		} else {
			borrow = 0
		}
		result[i] = (ad - bd) + '0'
	}

	int_str  := strings.clone(string(result[:max_int]))
	frac_str := strings.clone(string(result[max_int:]))
	delete(result)

	if len(int_str) == 0 {
		delete(int_str)
		int_str = strings.clone("0")
	}

	num := normalize_owned(DDB_Number{
		sign            = true,
		integer_part    = int_str,
		fractional_part = frac_str,
		exponent        = 0,
	})

	if len(num.integer_part) + len(num.fractional_part) > MAX_DDB_PRECISION {
		delete(num.integer_part)
		delete(num.fractional_part)
		return {}, false
	}

	return num, true
}

// Format a DDB_Number for display
format_ddb_number :: proc(num: DDB_Number) -> string {
	// Normalize first
	norm := normalize_ddb_number(num)

	// Check if it's effectively an integer
	if len(norm.fractional_part) == 0 && norm.exponent >= 0 {
		builder := strings.builder_make()
		if !norm.sign {
			strings.write_string(&builder, "-")
		}
		strings.write_string(&builder, norm.integer_part)
		// Add trailing zeros for positive exponent
		for _ in 0..<norm.exponent {
			strings.write_string(&builder, "0")
		}
		return strings.to_string(builder)
	}

	// Otherwise use full representation
	return ddb_number_to_string(norm)
}

// Clones a ddb number type
clone_ddb_number :: proc(num: DDB_Number) -> DDB_Number {
	return DDB_Number{
		sign = num.sign,
		integer_part = strings.clone(num.integer_part),
		fractional_part = strings.clone(num.fractional_part),
		exponent = num.exponent,
	}
}

// Helper: encode_varint (you already have this in your codebase)
@(private="file")
encode_varint :: proc(buf: ^bytes.Buffer, value: u64) {
	v := value
	for {
		byte_val := u8(v & 0x7F)
		v >>= 7
		if v != 0 {
			byte_val |= 0x80
		}
		bytes.buffer_write_byte(buf, byte_val)
		if v == 0 {
			break
		}
	}
}

// Helper: decode_varint
@(private="file")
decode_varint :: proc(data: []byte) -> (value: u64, bytes_read: int) {
	shift: u64 = 0
	for i in 0..<len(data) {
		byte_val := data[i]
		value |= u64(byte_val & 0x7F) << shift
		bytes_read = i + 1
		if (byte_val & 0x80) == 0 {
			return
		}
		shift += 7
	}
	return 0, 0
}