import logging
from pathlib import Path
from typing import Any, Callable, Coroutine, cast, final, overload

from ai_parade._toolkit.conversion_error import ConversionError
from ai_parade._toolkit.converters.converter import ConversionResult, Converter
from ai_parade._toolkit.converters.select import get_converter
from ai_parade._toolkit.datasets.Dataset import Dataset
from ai_parade._toolkit.enums.formats import ExactModelFormat
from ai_parade._toolkit.get_weights import get_weights_name
from ai_parade._toolkit.metadata.models.final import ModelMetadata
from ai_parade._toolkit.run.ModelRunner import LoadedModelRunner
from ai_parade._toolkit.run.RemoteRunner import RemoteError
from ai_parade._toolkit.run.runner_selector import get_runner
from ai_parade._toolkit.verify import compute_statistics

logger = logging.getLogger(__name__)


class Convert:
	"""
	Extensible converter that can convert between any two formats using available converters.
	Use by calling the instance (it has defined `__call__()`).

	## Path extension points:
	- `_get_result_path_impl`: gets the path to store conversion result JSON

	## Extension points:
	- `on_start`: called as first thing ever.
	- other extension points on each conversion (in other words even on conversions into intermediate formats):
	- `before_conversion`: called before each conversion step
	- `on_loaded_runner`: called when a runner is loaded, still before conversion is done
	- `save_converted_model`: saves converted model to appropriate location, if needed
	- `save_conversion_results`: saves multiple conversion results,
		should not be needed, internally calls `save_converted_model`


	"""

	def __init__(
		self,
		source_format: ExactModelFormat,
		target_format: ExactModelFormat,
		*,
		calibration_dataset: Dataset | None,
		verify_dataset: Dataset | None,
		override: bool,
		use_remote: bool,
		error_immediately: bool,
		metadata: ModelMetadata,
		output: Path | None = None,
		artifacts: Path | None = None,
		default_model: Path | None = None,
	):
		self.source_format = source_format
		self.target_format = target_format
		self.calibration_dataset = calibration_dataset
		self.verify_dataset = verify_dataset
		self.override = override
		self.use_remote = use_remote
		self.error_immediately = error_immediately
		self.metadata = metadata
		self.output = output.resolve() if output is not None else None
		self.artifacts = artifacts.resolve() if artifacts is not None else None

		if default_model is None:
			assert self.output is not None
			self.default_model = self.output / get_weights_name(source_format)
		else:
			self.default_model = default_model

	async def on_start(self, model_path: Path):
		"""
		Called when the converter starts. Before anything else. This method is always called once exactly.

		Arguments:
			model_path: Path to the source model, either weights or conversion result JSON.
		"""
		pass

	async def before_conversion(self):
		"""
		Called before each conversion step.
		"""
		pass

	async def on_loaded_runner(self, runner: LoadedModelRunner):
		"""
		To set up runner.
		Called when a runner is loaded, before each conversion begins. After before_conversion.

		Run the super implementation to set the output directory.
		"""
		assert self.artifacts is not None
		await runner.set_output_directory(self.artifacts)

	@overload
	def get_result_path(self, format_or_conversion_result: ExactModelFormat) -> Path:
		"""
		Gets the path to the conversion result JSON for the given format and hardware acceleration.
		"""
		pass

	@overload
	def get_result_path(self, format_or_conversion_result: ConversionResult) -> Path:
		"""
		Gets the path to the conversion result JSON for the given conversion result.
		"""
		pass

	@final
	def get_result_path(
		self, format_or_conversion_result: ExactModelFormat | ConversionResult
	):
		if isinstance(format_or_conversion_result, ConversionResult):
			format = format_or_conversion_result.format
		else:
			format = format_or_conversion_result
		return self._get_result_path_impl(format)

	def _get_result_path_impl(self, format: ExactModelFormat) -> Path:
		"""
		Gets the path to the conversion result JSON for the given format and hardware acceleration.
		"""

		assert self.output is not None
		model_path = self.output / get_weights_name(format)

		return model_path.with_name(model_path.name + ".json")

	def save_converted_model(self, conversion_result: ConversionResult) -> Path | None:
		"""
		Saves the converted model to the appropriate location and returns the new model path.

		If nothing is done, the model remains at its current path (in artifacts directory).
		The returned path is used in conversion results.
		"""
		new_model_path = self.get_result_path(conversion_result).with_suffix("")
		conversion_result.path.rename(new_model_path)
		return new_model_path

	async def save_failed_conversion_result(self):
		"""
		Saves the failed conversion result to the appropriate location.
		"""
		conversion_result = ConversionResult(
			status="failed",
			format=self.target_format,
			image_input=self.metadata.image_input,
			path=Path(),
			output="",
		)

		with open(self.get_result_path(conversion_result), "w", encoding="utf-8") as f:
			f.write(conversion_result.model_dump_json(indent=2))

	async def save_conversion_results(
		self,
		conversion_results: list[ConversionResult],
		current_target_format: ExactModelFormat | None,
	):
		"""
		Saves multiple conversion results and returns the one matching the current target format and hardware acceleration.
		"""
		requested_result = None
		for conversion_result in conversion_results:
			assert conversion_result.path.is_absolute()
			assert conversion_result.path.exists()
			new_model_path = (
				self.save_converted_model(conversion_result) or conversion_result.path
			)

			# update conversion result path
			result_path = self.get_result_path(conversion_result)
			conversion_result.path = (new_model_path).relative_to(
				result_path.parent, walk_up=True
			)

			# save conversion result
			with open(result_path, "w", encoding="utf-8") as f:
				f.write(conversion_result.model_dump_json(indent=2))
			logger.debug(
				f"Successfully converted model to {conversion_result.format} (at {new_model_path})"
			)

			# check if matches requested format
			if conversion_result.format == current_target_format:
				requested_result = conversion_result
		assert requested_result is not None
		return requested_result

	async def __call__(self, model_path: Path | None = None):
		model_path = model_path or self.default_model
		logger.info(
			f"Going to convert model `{self.metadata.name}` to {self.target_format}"
		)
		await self.on_start(model_path)

		assert model_path.exists(), f"Model path {model_path} does not exist"

		if model_path.suffix == ".json":
			self.metadata, model_path = self.metadata.with_conversion_results(
				model_path
			)

		converter_chains = get_converter(self.source_format, self.target_format)
		logger.debug(f"Converter chains: {converter_chains}")
		if not converter_chains or len(converter_chains) == 0:
			raise ConversionError(
				self.source_format, self.target_format, "Failed to find converter"
			)

		if self.verify_dataset is not None:
			return self.verify(
				lambda: self.__convert__(converter_chains, model_path), model_path
			)
		else:
			await self.__convert__(converter_chains, model_path)

	async def __convert__(
		self, converter_chains: list[list[Converter]], model_path: Path
	) -> ConversionResult | None:
		for converter_chain in converter_chains:
			# reset for each attempt
			current_source_weights_path = model_path
			current_metadata = self.metadata
			requested_result: ConversionResult | None = None
			try:
				for converter in converter_chain:
					# get current conversion variables
					current_source_format, current_target_format = converter.conversion
					current_target_result_path = self.get_result_path(
						current_target_format
					)
					if current_target_result_path.exists() and not self.override:
						# update for next iteration
						current_metadata, current_source_weights_path = (
							current_metadata.with_conversion_results(
								current_target_result_path
							)
						)
						logger.debug(
							f"Model already converted to {current_target_format}. Skipping"
						)
						continue

					# now lets convert
					logger.info(
						f"Converting from {converter.conversion[0]} to {converter.conversion[1]}"
					)
					await self.before_conversion()

					logger.info("Initializing runner")
					async with get_runner(
						current_metadata,
						current_source_weights_path,
						use_remote=self.use_remote,
					) as runner:
						await self.on_loaded_runner(runner)
						calibration_dataset = self.calibration_dataset
						if calibration_dataset is not None:
							calibration_dataset = calibration_dataset.preprocess(
								current_metadata.image_input
							)
						conversion_results = await runner.run(
							converter.create(
								Converter.Options(calibration_data=calibration_dataset)
							)
						)
					if not isinstance(conversion_results, list):
						conversion_results = [conversion_results]

					requested_result = await self.save_conversion_results(
						conversion_results, current_target_format
					)

					# update weights in next iteration
					current_metadata, current_source_weights_path = (
						current_metadata.with_conversion_results(
							self.get_result_path(requested_result), requested_result
						)
					)
				return requested_result
			except MemoryError as e:
				await self.save_failed_conversion_result()
				raise e
			except Exception as e:
				logger.error(
					f"Failed to convert model: {e if type(e) is RemoteError else str(e)}"
				)
				await self.save_failed_conversion_result()
				if self.error_immediately:
					raise e
				break
		raise ConversionError(
			self.source_format,
			self.target_format,
			f"No conversion lead to the target format, tried: {converter_chains}",
		)

	async def verify(
		self,
		convert: Callable[[], Coroutine[Any, Any, ConversionResult | None]],
		model_path: Path | None,
	) -> dict[str, float]:
		"""
		Run inference on both original and converted model and compare results.
		"""
		assert self.artifacts is not None
		assert self.verify_dataset is not None
		async with get_runner(
			self.metadata, model_path or self.default_model, use_remote=self.use_remote
		) as runner:
			await runner.set_output_directory(self.artifacts)
			base_results = await runner.inference(
				self.verify_dataset.preprocess(self.metadata.image_input)
			)

		result = await convert()
		if result is None:
			logger.warning("Conversion is cached, skipping verification")
			return {}

		async with get_runner(
			self.metadata, result.path, use_remote=self.use_remote
		) as runner:
			await runner.set_output_directory(self.artifacts)
			converted_results = await runner.inference(
				self.verify_dataset.preprocess(result.image_input)
			)
		return compute_statistics(cast(list, base_results), converted_results)
