import shutil
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Literal, override

from ai_parade._toolkit.converters._pytorch_base import PyTorchConvertBase
from ai_parade._toolkit.enums import ExactModelFormat, ModelFormat, Quantization
from ai_parade._toolkit.run import ModelRunner
from ai_parade._toolkit.run.ONNX import ONNXRunner

from .converter import ConversionResult, Converter

if TYPE_CHECKING:
	import torch  # type: ignore  # noqa: F401


def _create_ncnn_model_dir(output_weights_path: Path):
	output_weights_path.mkdir(exist_ok=True)
	(output_weights_path.parent / "model.ncnn.bin").rename(
		output_weights_path / "model.bin"
	)
	(output_weights_path.parent / "model.ncnn.param").rename(
		output_weights_path / "model.param"
	)


class PNNXPytorchConverter(Converter):
	"""
	https://github.com/Tencent/ncnn/tree/master/tools/pnnx/python
	https://github.com/pnnx/pnnx
	"""

	def __init__(
		self, quantization: Literal[Quantization.none] | Literal[Quantization.float16]
	):
		self.quantization = quantization

	@property
	@override
	def conversion(self):
		return (
			ExactModelFormat(ModelFormat.PyTorch),
			ExactModelFormat(ModelFormat.NCNN, self.quantization),
		)

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

	class PNNXConverter(PyTorchConvertBase):
		@override
		def _export_call(self, params: PyTorchConvertBase.Params):
			model, sample_inputs, output_weights_path, *_ = params

			import pnnx  # type: ignore

			_ = pnnx.export(
				model,
				f"{output_weights_path.parent}/model.pt",
				sample_inputs,
				fp16=self.conversion[1].quantization == Quantization.float16,
			)
			_create_ncnn_model_dir(output_weights_path)


class PNNXConverter(Converter):
	def __init__(self, conversion: tuple[ExactModelFormat, ExactModelFormat]):
		self._conversion = conversion

	@property
	@override
	def conversion(self):
		return self._conversion

	@override
	def create(
		self, options: Converter.Options
	) -> Callable[[ModelRunner], ConversionResult]:
		super().create(options)
		return self.convert

	def convert(self, runner: ModelRunner):
		assert isinstance(runner, ONNXRunner)
		import pnnx  # type: ignore
		import torch  # type: ignore

		shutil.copy(runner.weights, runner.output_directory / "model.onnx")
		# I guess it does not have an option to set output directory (for everything)
		_ = pnnx.convert(
			str(runner.output_directory / "model.onnx"),
			torch.rand(runner.model_metadata.image_input.shape),
			fp16=self.conversion[1].quantization != Quantization.float32,
		)
		_create_ncnn_model_dir(self.get_output_weights_path(runner, None))

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


capabilities: list[Converter] = [
	PNNXPytorchConverter(Quantization.none),
	PNNXPytorchConverter(Quantization.float16),
	PNNXConverter(
		(
			ExactModelFormat(ModelFormat.ONNX),
			ExactModelFormat(ModelFormat.NCNN),
		)
	),
	PNNXConverter(
		(
			ExactModelFormat(ModelFormat.ONNX),
			ExactModelFormat(ModelFormat.NCNN, Quantization.float16),
		)
	),
]
