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
}

// ============================================================================
// Arithmetic Operations
// ============================================================================

// Add two DDB_Numbers
add_ddb_numbers :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	// Convert to f64 for arithmetic (loses precision but matches current behavior)
	// TODO: Implement proper decimal arithmetic
	a_f64, a_ok := ddb_number_to_f64(a)
	b_f64, b_ok := ddb_number_to_f64(b)
	if !a_ok || !b_ok {
		return {}, false
	}

	result := a_f64 + b_f64
	return f64_to_ddb_number(result)
}

// Subtract two DDB_Numbers
subtract_ddb_numbers :: proc(a: DDB_Number, b: DDB_Number) -> (DDB_Number, bool) {
	// Convert to f64 for arithmetic
	a_f64, a_ok := ddb_number_to_f64(a)
	b_f64, b_ok := ddb_number_to_f64(b)
	if !a_ok || !b_ok {
		return {}, false
	}

	result := a_f64 - b_f64
	return f64_to_ddb_number(result)
}

// ============================================================================
// Conversion Helpers
// ============================================================================

// Convert DDB_Number to f64 (may lose precision)
ddb_number_to_f64 :: proc(num: DDB_Number) -> (f64, bool) {
	str := ddb_number_to_string(num)
	return strconv.parse_f64(str)
}

// Convert f64 to DDB_Number
f64_to_ddb_number :: proc(val: f64) -> (DDB_Number, bool) {
	// Format with enough precision to preserve the value
	str := fmt.aprintf("%.17g", val)
	return parse_ddb_number(str)
}

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