from pathlib import Path

import pytest
from typer.testing import CliRunner

from ai_parade._cli.ai_parade.main import app
from ai_parade._toolkit.converters.converter import ConversionResult

from ._common import files_equal, setup_cwd_any

runner = CliRunner()


class TestConvertCompare:
	repo_path = Path("fixtures/FooBarModel")

	@pytest.fixture(autouse=True)
	def copy_repo(self, tmp_path, monkeypatch):
		setup_cwd_any(
			tmp_path,
			monkeypatch,
			repo_path=self.repo_path,
			skip=["models"],
		)

	@pytest.mark.parametrize(
		["args", "file_names"],
		[
			pytest.param(*p, id=p[1])
			for p in [
				# LiteRT
				(
					"convert weights.pth FooBarModel tflite --artifacts build --separate-env",
					"tflite.float32.all.tflite",
				),
				(
					"convert weights.pth FooBarModel tflite --quantization float16 --separate-env",
					"tflite.float16.all.tflite",
				),
				(
					"convert weights.pth FooBarModel tflite --quantization int8 --calibration-dataset Random --separate-env",
					"tflite.int8.all.tflite",
				),
				# ncnn
				(
					"convert weights.pth FooBarModel ncnn --separate-env",
					"ncnn.float32.all.ncnn",
				),
				(
					"convert weights.pth FooBarModel ncnn --quantization float16 --separate-env",
					"ncnn.float16.all.ncnn",
				),
				(
					"convert weights.pth FooBarModel ncnn --quantization int8 --calibration-dataset Random --separate-env",
					"ncnn.int8.all.ncnn",
				),
				# Executorch
				(
					"convert weights.pth FooBarModel Executorch --hw-acceleration XNNPACK --separate-env",
					"executorch.float32.XNNPACK.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization float16 --hw-acceleration XNNPACK --separate-env",
					"executorch.float16.XNNPACK.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization int8 --calibration-dataset Random --hw-acceleration XNNPACK --separate-env",
					"executorch.int8.XNNPACK.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --hw-acceleration Vulkan --separate-env",
					"executorch.float32.Vulkan.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization float16 --hw-acceleration Vulkan --separate-env",
					"executorch.float16.Vulkan.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization int8 --calibration-dataset Random --hw-acceleration Vulkan --separate-env",
					"executorch.int8.Vulkan.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --hw-acceleration QNN --separate-env",
					"executorch.float32.QNN.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization float16 --hw-acceleration QNN --separate-env",
					"executorch.float16.QNN.pte",
				),
				(
					"convert weights.pth FooBarModel Executorch --quantization int8 --calibration-dataset Random --hw-acceleration QNN --separate-env",
					"executorch.int8.QNN.pte",
				),
				# ONNX
				(
					"convert weights.pth FooBarModel ONNX --separate-env",
					"onnx.float32.all.onnx",
				),
				(
					"convert weights.pth FooBarModel ONNX --quantization float16 --separate-env",
					"onnx.float16.all.onnx",
				),
				(
					"convert weights.pth FooBarModel ONNX --quantization int8 --calibration-dataset Random --separate-env",
					"onnx.int8.all.onnx",
				),
			]
		],
	)
	def test_conversion(self, args, file_names):
		result = runner.invoke(app, args)

		assert result.exit_code == 0

		for file_name in file_names if isinstance(file_names, list) else [file_names]:
			# check model
			expected_model: Path = (
				Path(__file__).parent.parent / self.repo_path / "models" / file_name
			)
			actual_model = Path(file_name)

			files_equal(expected_model, actual_model)

			# check json conversion results
			expected_result = expected_model.with_name(expected_model.name + ".json")
			actual_result = actual_model.with_name(actual_model.name + ".json")

			actual_result = ConversionResult.model_validate_json(
				actual_result.read_text()
			)
			expected_result = ConversionResult.model_validate_json(
				expected_result.read_text()
			)

			assert actual_result == expected_result
