from typing import override

from ai_parade._toolkit.enums import ModelFormat
from ai_parade._toolkit.run.ModelRunner import ModelRunner
from ai_parade._toolkit.run.ONNX import ONNXRunner

from .converter import ConversionResult, Converter, FormatConverter


class ONNX_TensorFlowConverter(FormatConverter):
	"""
	https://github.com/onnx/onnx-tensorflow/tree/main
	"""

	@property
	@override
	def format_conversion(self):
		return (ModelFormat.ONNX, ModelFormat.SavedModel)

	@override
	def create(self, options: Converter.Options):
		super().create(options)
		return self.convert

	def convert(self, runner: ModelRunner):
		assert isinstance(runner, ONNXRunner)
		from onnx_tf.backend import prepare  # type: ignore

		tf_rep = prepare(runner.model)  # prepare tf representation
		output_model_path = self.get_output_weights_path(runner, None)
		tf_rep.export_graph(output_model_path)  # export the model

		return ConversionResult(
			format=self.conversion[1],
			image_input=runner.model_metadata.image_input,
			path=output_model_path,
			output=runner.model_metadata.output,
		)
