from abc import ABC, abstractmethod
from pathlib import Path
from typing import Callable, NamedTuple, final, override

from ai_parade._toolkit.datasets.Dataset import Dataset
from ai_parade._toolkit.enums import (
	ExactModelFormat,
	ModelFormat,
)
from ai_parade._toolkit.get_weights import get_weights_name
from ai_parade._toolkit.run.ModelRunner import ModelRunner

from .ConversionResult import ConversionResult

Conversion = tuple[ExactModelFormat, ExactModelFormat]
"""
A conversion is a tuple of (input_format, output_format)
"""


class Converter(ABC):
	"""
	Converts a model from one format to another (including quantization).

	"""

	class Options(NamedTuple):
		"""
		Options for a converter.
		"""

		calibration_data: Dataset | None = None
		"""
		Dataset to use for quantization calibration, unused if no quantization
		"""

	@property
	@abstractmethod
	def conversion(self) -> Conversion:
		"""
		Available conversions in form of (input_format, output_format)
		"""
		pass

	@abstractmethod
	def create(
		self, options: Options
	) -> Callable[[ModelRunner], ConversionResult | list[ConversionResult]]:
		"""
		Creates a converter callable.

		The callable takes a ModelRunner as input and does not return anything.
		The converted model is stored in ModelRunner.output_directory as a `model.suffix` file(s).

		Implementation notes:
		- Use `get_output_weights_path` to get the path for the converted model

		Args:
			 calibration_data: dataset to use for quantization calibration, unused if no quantization

		Returns:
			 Callable that takes a ModelRunner as input and converts the model with it.
			 It returns the path to the converted model inside the runner's output directory
		"""

	def get_output_weights_path(
		self, runner: ModelRunner, options: Options | None
	) -> Path:
		return runner.output_directory / get_weights_name(self.conversion[1])


class FormatConverter(Converter):
	"""
	Subclass of Converter that converts model formats without quantization.
	"""

	@property
	@final
	@override
	def conversion(self) -> Conversion:
		formats = self.format_conversion
		return (
			ExactModelFormat(formats[0]),
			ExactModelFormat(formats[1]),
		)

	@property
	@abstractmethod
	def format_conversion(self) -> tuple[ModelFormat, ModelFormat]:
		pass
