use math;
use types;

// premature optimization ftw!

export def DIRECT: u8 = 0;
export def CONSTANT: u8 = 1;
export def INLINE_MAX: u8 = ((size(*opaque) * 2 - 1) / 2): u8;

export type Array = union {
	// (2024 note): this used to use duplicate struct member names, which i guess
	// was a compiler bug and doesn't work anymore.
	struct { // direct
		palette_size: u8, // == DIRECT
		data_direct: *[*]u16,
	},
	struct { // constant
		palette_size_const: u8, // == CONSTANT
		data_constant: u16,
	},
	struct { // indirect inline
		palette_size_indirect_inline: u8, // <= INLINE_MAX, != DIRECT, != CONSTANT
		palette_inline: [INLINE_MAX]u16,
		data_indirect: *[*]u8,
	},
	struct { // indirect ptr
		palette_size_indirect_ptr: u8, // > INLINE_MAX
		palette_ptr: *[*]u16,
		// data_indirect: *[*]u8,
		// (though it would result in an incorrect offset... yeah,
		// i don't know what i was thinking here at all.)
	},
};

export fn new(palette_size: u8, arraylen: size) Array = {
	let array = Array {
		palette_size = palette_size,
	};

	if (palette_size == DIRECT) {
		array.data_direct =
			alloc([]: [0]u16, arraylen): *[*]u16;
	} else if (palette_size != CONSTANT) {
		array.data_indirect = alloc([]: [0]u8,
			indirect_size(palette_size, arraylen)): *[*]u8;
		if (palette_size > INLINE_MAX) {
			array.palette_ptr =
				alloc([]: [0]u16, palette_size): *[*]u16;
		};
	};

	return array;
};

export fn clear(array: *Array, fill_value: u16) void = {
	if (array.palette_size == DIRECT) {
		free(array.data_direct);
	} else if (array.palette_size != CONSTANT) {
		free(array.data_indirect);
		if (array.palette_size > INLINE_MAX)
			free(array.palette_ptr);
	};

	array.palette_size = CONSTANT;
	array.data_constant = fill_value;
};

// XXX: should be @inline
export fn get(array: *Array, i: size) u16 = {
	if (array.palette_size == DIRECT) {
		return array.data_direct[i];
	} else if (array.palette_size == CONSTANT) {
		return 0;
	} else {
		const log2nbits = palette_log2nbits(array.palette_size);
		const nbits = 1 << log2nbits;
		const valuemask = (1 << nbits) - 1;
		const log2nperbyte = 3 - log2nbits;
		const nperbyte = 1 << log2nperbyte;
		const indexlow = nperbyte - 1;

		const byte = array.data_indirect[i >> log2nperbyte];
		const shift = (i: u8 & indexlow) << log2nbits;
		return (byte >> shift) & valuemask;
	};
};

// XXX: should be @inline
export fn set(array: *Array, i: size, ref: u16) void = {
	if (array.palette_size == DIRECT) {
		array.data_direct[i] = ref;
	} else if (array.palette_size == CONSTANT) {
		assert(ref == 0);
	} else {
		const log2nbits = palette_log2nbits(array.palette_size);
		const nbits = 1 << log2nbits;
		const valuemask = (1 << nbits) - 1;
		const log2nperbyte = 3 - log2nbits;
		const nperbyte = 1 << log2nperbyte;
		const indexlow = nperbyte - 1;

		const byte = &array.data_indirect[i >> log2nperbyte];
		const shift = (i: u8 & indexlow) << log2nbits;
		*byte = *byte & ~(valuemask << shift) | ref: u8 << shift;
	};
};

// XXX: should be @inline
export fn get_palette(array: *Array) []u16 = {
	if (array.palette_size == DIRECT) {
		abort();
	} else if (array.palette_size == CONSTANT) {
		return &array.data_constant: *[1]u16;
	} else if (array.palette_size <= INLINE_MAX) {
		return array.palette_inline[..array.palette_size];
	} else {
		return array.palette_ptr[..array.palette_size];
	};
};

// XXX: should be @inline
export fn lookup(array: *Array, ref: u16) u16 = {
	if (array.palette_size == DIRECT) {
		return ref;
	};
	return get_palette(array)[ref];
};

// XXX: should be @inline
export fn indirect_size(palette_size: u8, arraylen: size) size =
	arraylen << palette_log2nbits(palette_size) >> 3;

// XXX: should be @inline
export fn palette_log2nbits(palette_size: u8) u8 = {
	assert(palette_size != DIRECT);
	const nbits = math::bit_size_u8(palette_size - 1);
	return math::bit_size_u8(nbits - 1);
};

export fn grow_palette(array: *Array, arraylen: size, newsize: u8) void = {
	if (array.palette_size == newsize) return;
	assert(
		(array.palette_size - 1) < (newsize - 1),
		"new size must be >= old size"
	);

	let newarray = *array;
	newarray.palette_size = newsize;

	defer *array = newarray;

	// By this point we know that the palette has grown in size, so it
	// always needs to be reallocated unless the new size is DIRECT,
	// in which case there will be no palette.

	const palette = get_palette(array);

	if (newsize == DIRECT) {
		void;
	} else if (newsize <= INLINE_MAX) {
		newarray.palette_inline[..len(palette)] = palette;
	} else {
		let newpalette = alloc(palette, newsize);
		newarray.palette_ptr = newpalette: *[*]u16;
	};

	defer if (array.palette_size > INLINE_MAX) {
		free(array.palette_ptr);
	};

	if (array.palette_size == CONSTANT) {
		if (newsize == DIRECT) {
			newarray.data_direct = alloc([array.data_constant...],
				arraylen): *[*]u16;
		} else {
			newarray.data_indirect = alloc([0...],
				indirect_size(newsize, arraylen)): *[*]u8;
		};

		return;
	};

	if (newsize == DIRECT) {
		newarray.data_direct = alloc([]: [0]u16, arraylen): *[*]u16;

		for (let i = 0z; i < arraylen; i += 1) {
			newarray.data_direct[i] = palette[get(array, i)];
		};

		free(array.data_indirect);
		return;
	};

	// We need to reallocate the data array if the number of bits per item
	// increased.

	const log2nbits = palette_log2nbits(array.palette_size);
	const log2newnbits = palette_log2nbits(newsize);

	if (log2nbits != log2newnbits) {
		newarray.data_indirect = alloc([]: [0]u8,
			indirect_size(newsize, arraylen)): *[*]u8;

		for (let i = 0z; i < arraylen; i += 1) {
			set(&newarray, i, get(array, i));
		};

		free(array.data_indirect);
	};
};