import logging
import traceback
from abc import ABC, abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING, NamedTuple

from ai_parade._toolkit.data.ImageInput import ImageInput
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 ModelRunner
from ai_parade._toolkit.run.PyTorch import PyTorchRunner

from .converter import Conversion, ConversionResult, Converter

if TYPE_CHECKING:
	import torch  # type: ignore

logger = logging.getLogger(__name__)


class PyTorchConvertBase(ABC):
	def __init__(self, conversion: Conversion, options: Converter.Options):
		self.conversion = conversion
		self.options = options

	class Params(NamedTuple):
		model: "torch.nn.Module"
		sample_inputs: tuple["torch.Tensor"]
		output_weights_path: Path
		default_results: ConversionResult
		metadata: ModelMetadata

	@abstractmethod
	def _export_call(
		self, params: Params
	) -> ConversionResult | list[ConversionResult] | None:
		pass

	def __call__(
		self, runner: ModelRunner
	) -> ConversionResult | list[ConversionResult]:
		import torch  # type: ignore

		assert isinstance(runner, PyTorchRunner)

		match runner.model_metadata.image_input.dataType:
			case ImageInput.DataType.float32:
				dtype = torch.float32
			case ImageInput.DataType.float16:
				dtype = torch.float16
			case ImageInput.DataType.uint8:
				dtype = torch.uint8

		sample_inputs = torch.rand(runner.model_metadata.image_input.shape, dtype=dtype)
		runner.model.eval()

		def run(
			sample_inputs: "torch.Tensor",
		) -> ConversionResult | list[ConversionResult]:
			# run the model to trigger potential lazy initialization
			runner.model(sample_inputs)

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

			results = self._export_call(
				PyTorchConvertBase.Params(
					runner.model,
					(sample_inputs,),
					self.get_output_weights_path(runner),
					default_results=default_results,
					metadata=runner.model_metadata,
				)
			)

			return results or default_results

		try:
			runner.model.to("cpu")
			return run(sample_inputs.to("cpu"))
		except Exception as e:
			msg = str(e)
			if "cpu" in msg or "gpu" in msg or "device" in msg:
				logger.warning("Conversion on CPU model failed, trying cuda instead.")
				logger.info(f"Failed with error: '{msg}'")
				logger.debug(traceback.format_exc())
				runner.model.to("cuda")
				return run(sample_inputs.to("cuda"))
			else:
				raise e

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