from pathlib import Path

from ai_parade._toolkit.enums import ExactModelFormat, format_suffixes
from ai_parade._toolkit.runtime_metadata import ModelAttributes, format_distinguishes


def get_weights_name(format: ExactModelFormat) -> Path:
	"""
	Get path to model weights, this path may not exist

	Example:
	>>> get_weights_name(ExactModelFormat(ModelFormat.ONNX, Quantization.none))
	'onnx.all.all.onnx'

	>>> get_weights_name(ExactModelFormat(ModelFormat.ExecutorTorch, Quantization.int8, HardwareAcceleration.XNNPACK))
	'executorch.int8.xnnpack.pte'
	"""
	if ModelAttributes.HWAcceleration in format_distinguishes[format.format]:
		assert format.hw_acceleration_restriction is not None
		hw_acceleration = format.hw_acceleration_restriction.value
	else:
		hw_acceleration = "all"
	name = ".".join(
		[
			format.format.value,
			(
				format.quantization.value
				if ModelAttributes.Quantization in format_distinguishes[format.format]
				else "all"
			),
			hw_acceleration,
		]
	)
	return Path(name + format_suffixes[format.format])
