from dataclasses import dataclass
from enum import StrEnum

from .AIRuntime import AIRuntime
from .hw_acceleration import HardwareAcceleration


class ModelFormat(StrEnum):
	PyTorch = "pytorch"
	ExecutorTorch = "executorch"
	SavedModel = "savedmodel"
	LiteRT = "tflite"
	ONNX = "onnx"
	SNPE = "snpe"
	NCNN = "ncnn"
	MNN = "mnn"
	bolt = "bolt"
	MACE = "mace"
	PaddleLite = "paddlelite"
	# New AI Runtime: update here


class Quantization(StrEnum):
	none = "float32"
	float32 = "float32"
	float16 = "float16"
	int8 = "int8"
	int4 = "int4"


@dataclass(frozen=True)
class ExactModelFormat:
	"""
	Represents a exact model format.
	"""

	format: ModelFormat
	"""
	Model format
	"""
	quantization: Quantization = Quantization.none
	"""
	Quantization of the model
	"""
	hw_acceleration_restriction: HardwareAcceleration | None = None
	"""
	Hardware acceleration of the converted model
	"""
	chip_compatibility: tuple[str, ...] | None = None
	"""
	Chip compatibility of the converted model
	"""


formats_preferred_runtime: dict[ModelFormat, AIRuntime] = {
	ModelFormat.LiteRT: AIRuntime.LiteRT,
	ModelFormat.NCNN: AIRuntime.ncnn,
	ModelFormat.ONNX: AIRuntime.ONNXRuntime,
	ModelFormat.PyTorch: AIRuntime.PyTorch,
	ModelFormat.SavedModel: AIRuntime.TensorFlow,
}

format_suffixes: dict[ModelFormat, str] = {
	ModelFormat.SavedModel: ".savedmodel",
	ModelFormat.PyTorch: ".pth",
	ModelFormat.SavedModel: ".pb",
	ModelFormat.LiteRT: ".tflite",
	ModelFormat.ONNX: ".onnx",
	# this should be directory with ncnn.param and ncnn.bin
	ModelFormat.NCNN: ".ncnn",
	ModelFormat.ExecutorTorch: ".pte",
	# New AI Runtime: update here
}
