import json
from pathlib import Path
from typing import Any, Callable, Iterable

import numpy as np

from ai_parade._toolkit.data.output import ModelOutput
from ai_parade._toolkit.datasets.Dataset import Dataset
from ai_parade._toolkit.metadata.models.final import ModelMetadata
from ai_parade._toolkit.run.runner_selector import get_runner


async def inference(
	weights_path: Path,
	model_metadata: ModelMetadata,
	dataset: Dataset,
	# common
	artifacts: Path,
	output: Path | None,
	metrics: dict[str, Callable[[Iterable[ModelOutput], Iterable[ModelOutput]], Any]],
	use_remote: bool,
):
	"""
	Run inference on a model using a selected dataset and print results.

	Args:
		weights_path: Path to weights or a conversion result JSON
		model_metadata: ModelMetadata
		dataset: Dataset without preprocessing (preprocessing done in this method)
		artifacts: Path
		output: Path path to json file to store results
		metrics: dict[str, Callable[[Iterable[ModelOutput], Iterable[ModelOutput]], Any]]
		use_remote: bool
	"""
	assert output is None or output.suffix == ".json"

	if weights_path.suffix == ".json":
		model_metadata, weights_path = model_metadata.with_conversion_results(
			weights_path
		)
	dataset = dataset.preprocess(model_metadata.image_input)

	async with get_runner(
		model_metadata, weights_path, use_remote=use_remote
	) as runner:
		await runner.set_output_directory(artifacts)
		raw_results = await runner.inference(dataset)

	gt = [label for _, label, _ in dataset.no_batch_iterator()]

	results = {}
	if output is not None:
		for metric, metric_fx in metrics.items():
			results[metric] = metric_fx(raw_results, gt)

		class NumpyEncoder(json.JSONEncoder):
			def default(self, o):
				if isinstance(o, np.ndarray):
					return o.tolist()
				elif isinstance(o, np.generic):
					return o.item()
				return super().default(o)

		with open(output, "w") as f:
			json.dump(results, f, cls=NumpyEncoder)

	return results
