package encoding_base32

// @note(zh): Encoding utility for Base32
// A secondary param can be used to supply a custom alphabet to
// @link(encode) and a matching decoding table to @link(decode).
// If none is supplied it just uses the standard Base32 alphabet.
// Incase your specific version does not use padding, you may
// truncate it from the encoded output.

// Error represents errors that can occur during base32 decoding operations.
// See RFC 4648 sections 3.2, 4 and 6.
Error :: enum {
	None,
	Invalid_Character, // Input contains characters outside of base32 alphabet (A-Z, 2-7)
	Invalid_Length,    // Input length is not valid for base32 (must be a multiple of 8 with proper padding)
	Malformed_Input,   // Input has improper structure (wrong padding position or incomplete groups)
}

ENC_TABLE := [32]byte {
	'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
	'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
	'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
	'Y', 'Z', '2', '3', '4', '5', '6', '7',
}

PADDING :: '='

DEC_TABLE := [?]u8 {
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0, 26, 27, 28, 29, 30, 31,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
	15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,  0,  0,  0,  0,  0,
	 0,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
	15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
	 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
}

encode :: proc(data: []byte, ENC_TBL := ENC_TABLE, allocator := context.allocator) -> string {
	out_length := (len(data) + 4) / 5 * 8
	out := make([]byte, out_length)
	_encode(out, data)
	return string(out)
}

@private
_encode :: proc(out, data: []byte, ENC_TBL := ENC_TABLE, allocator := context.allocator) {
	out := out
	data := data

	for len(data) > 0 {
		carry: byte
		switch len(data) {
		case:
			out[7] = ENC_TABLE[data[4] & 0x1f]
			carry = data[4] >> 5
			fallthrough
		case 4:
			out[6] = ENC_TABLE[carry | (data[3] << 3) & 0x1f]
			out[5] = ENC_TABLE[(data[3] >> 2) & 0x1f]
			carry = data[3] >> 7
			fallthrough
		case 3:
			out[4] = ENC_TABLE[carry | (data[2] << 1) & 0x1f]
			carry = (data[2] >> 4) & 0x1f
			fallthrough
		case 2:
			out[3] = ENC_TABLE[carry | (data[1] << 4) & 0x1f]
			out[2] = ENC_TABLE[(data[1] >> 1) & 0x1f]
			carry = (data[1] >> 6) & 0x1f
			fallthrough
		case 1:
			out[1] = ENC_TABLE[carry | (data[0] << 2) & 0x1f]
			out[0] = ENC_TABLE[data[0] >> 3]
		}

		if len(data) < 5 {
			out[7] = byte(PADDING)
			if len(data) < 4 {
				out[6] = byte(PADDING)
				out[5] = byte(PADDING)
				if len(data) < 3 {
					out[4] = byte(PADDING)
					if len(data) < 2 {
						out[3] = byte(PADDING)
						out[2] = byte(PADDING)
					}
				}
			}
			break
		}
		data = data[5:]
		out = out[8:]
	}
}

decode :: proc(data: string, DEC_TBL := DEC_TABLE, allocator := context.allocator) -> ([]byte, Error) {
	if len(data) == 0 {
		return nil, .None
	}

	// Calculate maximum possible output size and allocate buffer
	out_len := (len(data) * 5 + 7) / 8 // Ceiling division to ensure enough space
	out := make([]byte, out_len, allocator)

	outi := 0
	data := data

	end := false
	for len(data) > 0 && !end {
		dbuf : [8]byte
		dlen := 8

		for j := 0; j < 8; {
			if len(data) == 0 {
				dlen, end = j, true
				break
			}
			input := data[0]
			data = data[1:]
			if input == byte(PADDING) && j >= 2 && len(data) < 8 {
				if len(data) + j < 8 - 1 {
					return nil, .Malformed_Input
				}
				for k := 0; k < 8-1-j; k += 1 {
					if len(data) < k || data[k] != byte(PADDING) {
						return nil, .Malformed_Input
					}
				}
				dlen, end = j, true
				if dlen == 1 || dlen == 3 || dlen == 6 {
					return nil, .Invalid_Length
				}
				break
			}

			decoded := DEC_TBL[input]
			if decoded == 0 && input != byte(ENC_TABLE[0]) {
				return nil, .Invalid_Character
			}
			dbuf[j] = decoded
			j += 1
		}

		// Ensure we have enough space in output buffer
		needed := 5  // Each full 8-char block produces 5 bytes
		if outi + needed > len(out) {
			return nil, .Invalid_Length
		}

		// Process complete input blocks
		switch dlen {
		case 8:
			if len(dbuf) < 8 { return nil, .Invalid_Length }
			out[outi + 4] = dbuf[6] << 5 | dbuf[7]
			fallthrough
		case 7:
			if len(dbuf) < 7 { return nil, .Invalid_Length }
			out[outi + 3] = dbuf[4] << 7 | dbuf[5] << 2 | dbuf[6] >> 3
			fallthrough
		case 5:
			if len(dbuf) < 5 { return nil, .Invalid_Length }
			out[outi + 2] = dbuf[3] << 4 | dbuf[4] >> 1
			fallthrough
		case 4:
			if len(dbuf) < 4 { return nil, .Invalid_Length }
			out[outi + 1] = dbuf[1] << 6 | dbuf[2] << 1 | dbuf[3] >> 4
			fallthrough
		case 2:
			if len(dbuf) < 2 { return nil, .Invalid_Length }
			out[outi + 0] = dbuf[0] << 3 | dbuf[1] >> 2
		}
		outi += 5
	}

	// Trim output buffer to actual size
	if outi < len(out) {
		out = out[:outi]
	}

	return out, .None
}