from pathlib import Path

from ai_parade._toolkit.enums import AIRuntime, formats_preferred_runtime
from ai_parade._toolkit.metadata.custom import Customizations
from ai_parade._toolkit.metadata.models.api import RunnerGetter
from ai_parade._toolkit.metadata.models.final import ModelMetadata

from .ModelRunner import ModelRunner
from .NotARunner import NotARunner
from .RemoteRunner import RemoteRunner


def get_runner(
	model_metadata: ModelMetadata,
	weights: Path,
	runtime: AIRuntime | None = None,
	use_remote: bool = True,
	customizations: Customizations | None = None,
	runner_name: str | None = None,
) -> ModelRunner:
	"""Returns appropriate runner for the model based on the parameters

	Args:
		model_metadata: metadata
		weights: Model weights
		runtime: Desired AI runtime
		use_remote: If False, don't use remote runner. Defaults to True.
		customizations: Customizations. Defaults to None.
		runner_name: Name of the runner in the customizations. Defaults to None, required if customizations is not None

	Returns:
		Initialized ModelRunner
	"""
	runtime = runtime or formats_preferred_runtime[model_metadata.format.format]
	if model_metadata.get_custom_runner is not None:
		assert formats_preferred_runtime[model_metadata.format.format] == runtime
		runner = model_metadata.get_custom_runner(model_metadata, weights)
		if not use_remote and isinstance(runner, RemoteRunner):
			return runner.local_runner(model_metadata, weights)
		return runner
	else:
		local_getter = _choose_local_runner(runtime, customizations, runner_name)
		if not use_remote:
			return local_getter(model_metadata, weights)

		return RemoteRunner(model_metadata, runtime, weights, local_getter)


def _choose_local_runner(
	runtime: AIRuntime,
	customizations: Customizations | None = None,
	runner_name: str | None = None,
) -> RunnerGetter:
	"""Chooses default runner or runner from customizations

	Args:
		model_format: Desired model format
		customizations: Customizations. Defaults to None.
		runner_name: Name of the runner in the customizations. Defaults to None, required if customizations is not None

	Returns:
		ModelRunner in form of RunnerGetter
	"""
	if customizations is not None and runner_name is not None:
		return customizations.runners[runner_name]

	else:
		match runtime:
			case AIRuntime.PyTorch:
				from .PyTorch import PyTorchRunner

				return PyTorchRunner
			case AIRuntime.ONNXRuntime:
				from .ONNX import ONNXRunner

				return ONNXRunner
			case AIRuntime.TensorFlow:
				from .TensorFlow import TensorFlowRunner

				return TensorFlowRunner
			case AIRuntime.ExecuTorch:
				return NotARunner
			case AIRuntime.LiteRT:
				return NotARunner
			case AIRuntime.ncnn:
				return NotARunner
			case _:
				raise RuntimeError(f"Unknown model format `{runtime}`")
			# New AI Runtime: register runner - ModelRunner derived class
