import type { Readable } from "@zarrita/storage";
import { create_codec_pipeline } from "../codecs.js";
import type { Location } from "../hierarchy.js";
import type { Chunk } from "../metadata.js";
import { assert, type ShardingCodecMetadata } from "../util.js";

const MAX_BIG_UINT = 18446744073709551615n;

export function create_sharded_chunk_getter<Store extends Readable>(
	location: Location<Store>,
	shard_shape: number[],
	encode_shard_key: (coord: number[]) => string,
	sharding_config: ShardingCodecMetadata["configuration"],
) {
	assert(location.store.getRange, "Store does not support range requests");
	let get_range = location.store.getRange.bind(location.store);
	let index_shape = shard_shape.map(
		(d, i) => d / sharding_config.chunk_shape[i],
	);
	let index_codec = create_codec_pipeline({
		data_type: "uint64",
		shape: [...index_shape, 2],
		codecs: sharding_config.index_codecs,
	});

	let checksum_size = 4;
	let index_size = 16 * index_shape.reduce((a, b) => a * b, 1);
	let cache: Record<string, Promise<Chunk<"uint64"> | null>> = {};
	return async (
		chunk_coord: number[],
		options?: Parameters<Store["get"]>[1],
	) => {
		let shard_coord = chunk_coord.map((d, i) => Math.floor(d / index_shape[i]));
		let shard_path = location.resolve(encode_shard_key(shard_coord)).path;

		if (!(shard_path in cache)) {
			cache[shard_path] = (async () => {
				let bytes = await get_range(
					shard_path,
					{
						suffixLength: index_size + checksum_size,
					},
					options,
				);
				return bytes ? await index_codec.decode(bytes) : null;
			})().catch((err) => {
				delete cache[shard_path];
				throw err;
			});
		}
		let index = await cache[shard_path];

		if (index === null) {
			return undefined;
		}

		let { data, shape, stride } = index;
		let linear_offset = chunk_coord
			.map((d, i) => d % shape[i])
			.reduce((acc, sel, idx) => acc + sel * stride[idx], 0);

		let offset = data[linear_offset];
		let length = data[linear_offset + 1];
		// write null chunk when 2^64-1 indicates fill value
		if (offset === MAX_BIG_UINT && length === MAX_BIG_UINT) {
			return undefined;
		}
		return get_range(
			shard_path,
			{
				offset: Number(offset),
				length: Number(length),
			},
			options,
		);
	};
}