import asyncio
import logging
import sys
from pathlib import Path
from pprint import pprint
from typing import Annotated

import rich
import rich.traceback
import typer

import ai_parade._compat  # noqa: F401
from ai_parade._toolkit.install.utilities import check_missing_dependencies
from ai_parade.datasets import DatasetSplit, ImageNet, ImageNet_VID, Random
from ai_parade.logging import TRACE_LVL
from ai_parade.toolkit import (
	Convert,
	Dataset,
	ExactModelFormat,
	HardwareAcceleration,
	ModelFormat,
	Quantization,
	RemoteError,
	get_local_metadata,
	get_runner,
)
from ai_parade.toolkit import (
	get_ai_parade_logger as get_toolkit_logger,
)
from ai_parade.toolkit import inference as run_inference

logger = logging.getLogger(__name__)

app = typer.Typer()

RepositoryPathType = Annotated[
	Path,
	typer.Option(
		"--repository",
		help="Path to repository with pyproject.toml",
		file_okay=False,
		exists=True,
		resolve_path=True,
	),
]
WeightsPathTypeArgDef = typer.Argument(
	help="Path to the model file (all in one file), or directory (if multiple files are needed), or json conversion result.",
	exists=True,
	resolve_path=True,
)
WeightsPathTypeOptional = Annotated[Path | None, WeightsPathTypeArgDef]
WeightsPathType = Annotated[Path, WeightsPathTypeArgDef]

ModelNameTypeArgDef = typer.Argument(
	help="Name of model (as in the [tool.ai-parade.name] property in pyproject.toml file)"
)
ModelNameType = Annotated[str, ModelNameTypeArgDef]
ModelNameTypeOptional = Annotated[str | None, ModelNameTypeArgDef]

ArtifactsDirType = Annotated[
	Path,
	typer.Option(
		help="Directory for artifacts (where logs, dumps etc. will be stored).",
		file_okay=False,
	),
]

OutputTypeOptionDef = typer.Option(
	"--output", "-o", help="Output directory for converted weights."
)

OutputType = Annotated[Path, OutputTypeOptionDef]
OutputTypeOptional = Annotated[Path | None, OutputTypeOptionDef]

SeparateEnvironmentType = Annotated[
	bool,
	typer.Option(
		"--separate-environment",
		"--separate-env",
		"--separate-venv",
		help="Separates environment (python interpreter and venv) from the local one. Setting this will complicate debugging.",
		show_default="False - local environment",
	),
]

NEED_CALIBRATION = [Quantization.int4, Quantization.int8]

verbosity_level = 0


def _set_verbosity(verbose: "VerboseType"):
	toolkit_logger = get_toolkit_logger()
	this_logger = logging.getLogger("ai_parade._cli.ai_parade")
	global verbosity_level
	verbosity_level = verbose.count(True)

	if verbosity_level < 2:
		logging.basicConfig(
			format="[%(levelname)-9s%(asctime)s] %(message)s",
			datefmt="%H:%M:%S",
		)
	else:
		logging.basicConfig(
			format="[%(levelname)-9s%(asctime)s] %(filename)s#%(lineno)d - %(message)s",
			datefmt="%H:%M:%S",
		)
	try:
		match verbosity_level:
			case 0:
				this_logger.setLevel(logging.INFO)
			case 1:
				toolkit_logger.setLevel(logging.INFO)
				this_logger.setLevel(logging.DEBUG)
			case 2:
				toolkit_logger.setLevel(logging.DEBUG)
				this_logger.setLevel(TRACE_LVL)
			case 3:
				toolkit_logger.setLevel(TRACE_LVL)
				this_logger.setLevel(TRACE_LVL)
			case _:
				logging.getLogger().setLevel(TRACE_LVL)
	except Exception as e:
		logging.basicConfig(level=logging.INFO)
		this_logger.error(f"Failed to set verbosity level: {e}")


VerboseType = Annotated[
	list[bool],
	typer.Option(
		"--verbose",
		"-v",
		help="Enable verbose logging, repeat for more verbosity (current max: `-vvvv`)",
		callback=_set_verbosity,
	),
]


def handle_remote_error(e: RemoteError):
	rich.print(e.rich_traceback)
	if verbosity_level >= 3:
		rich.print(rich.traceback.Traceback.from_exception(*sys.exc_info()))  # type: ignore
	raise typer.Exit(code=1)


def _get_dataset(
	dataset_name: str, dataset_path: Path | None, split: DatasetSplit
) -> Dataset:
	match dataset_name:
		case "ImageNet":
			if dataset_path is None:
				raise ValueError("Dataset path is required for ImageNet dataset")
			dataset = ImageNet(dataset_path, split)
		case "ImageNet_VID":
			if dataset_path is None:
				raise ValueError("Dataset path is required for ImageNet_VID dataset")
			dataset = ImageNet_VID(dataset_path, split)
		case "Random":
			dataset = Random(seed=442)
		case _:
			raise NotImplementedError(f"Dataset {dataset_name} not implemented")
	return dataset


@app.command()
def ls(
	repository: RepositoryPathType = Path("."),
	verbose: VerboseType = [],
):
	"""
	List all available models in the repository.
	"""
	metadata_list = asyncio.run(get_local_metadata(repository))
	print(
		"Models metadata parsed:\n\t" + "\n\t".join(m.name for m in metadata_list),
		flush=True,
	)


@app.command()
def check(
	model_name: ModelNameTypeOptional = None,
	weights_path: WeightsPathTypeOptional = None,
	skip_dependencies: Annotated[
		bool,
		typer.Option(help="If set to True, don't check dependencies."),
	] = False,
	# reinstall_dependencies: Annotated[
	# bool,
	# typer.Option(
	# help="If set to True, create a new environment and install dependencies for each model."
	# ),
	# ] = False,
	pin_installed: Annotated[
		bool,
		typer.Option(
			help=(
				"Set to True to pin dependencies to specific version into `requirements.txt`. "
				"Uses currently installed dependencies... its pip freeze equivalent"
			)
		),
	] = False,
	pin_defined: Annotated[
		bool,
		typer.Option(
			help=(
				"Set to True to pin dependencies to specific version into `requirements.txt`. "
				"If the repository contains a lock file they will be converted to requirements.txt"
				"Supported dependency sources: "
				"requirements.in, setup.cfg, setup.py, pyproject.toml, requirements.txt; "
				"Supported lock files: "
				"pipFile.lock, poetry.lock"
			)
		),
	] = False,
	repository: RepositoryPathType = Path("."),
	verbose: VerboseType = [],
):
	"""
	Check if model repository is correctly set up.

	It runs the following steps:

	1. Check if installed dependencies are pinned.
		- Skip with `--skip-dependencies`

	2. Parse metadata and select model (`MODEL_NAME`)

	3. Run inference on random data.
		- Only run if model name and weights are provided (`MODEL_NAME` and `WEIGHTS_PATH`)
	"""
	has_missing_dependencies = False
	if not skip_dependencies:
		has_missing_dependencies = asyncio.run(
			check_missing_dependencies(
				Path.cwd(), pin_installed=pin_installed, pin_defined=pin_defined
			)
		)
		logger.info("Dependencies checked!")

	metadata_list = asyncio.run(get_local_metadata(repository))
	logger.info("Model metadata successfully parsed")

	if model_name is not None:
		assert weights_path is not None
		metadata = next(
			x for x in metadata_list if x.name == model_name if model_name is not None
		)
		logger.info(f"Selected model: {metadata.name}")

		# if reinstall_dependencies:
		# with TemporaryDirectory() as venv:
		# raise NotImplementedError
		# runner.inference(Random(metadata.image_input))
		# else:
		async def run_inference():
			async with get_runner(metadata, weights_path) as runner:
				await runner.inference(Random(metadata.image_input).take(2))

		asyncio.run(run_inference())

	exit_code = 0
	if has_missing_dependencies:
		exit_code |= 8
	raise typer.Exit(code=exit_code)


@app.command()
def convert(
	weights_path: WeightsPathType,
	model_name: ModelNameType,
	model_format: Annotated[
		ModelFormat, typer.Argument(case_sensitive=False, help="Target model format")
	],
	quantization: Annotated[
		Quantization,
		typer.Option(
			"--quantization",
			"--quantize",
			case_sensitive=False,
			help="Quantization of the target model",
		),
	] = Quantization.none,
	hw_acceleration: Annotated[
		HardwareAcceleration,
		typer.Option(
			case_sensitive=False,
			help="Hardware acceleration of the target model. It is used only by Executorch",
		),
	] = HardwareAcceleration.CPU,
	override: Annotated[
		bool,
		typer.Option(
			help="Overwrite already converted weights if they exist (including intermediate ones)."
		),
	] = False,
	# calibration for quantization
	calibration_dataset: Annotated[
		str | None,
		typer.Option(
			"--calibration-dataset",
			help="Dataset name used to calibrate INT4/INT8 conversions. Required for INT4/INT8 conversions.",
		),
	] = None,
	calibration_dataset_path: Annotated[
		Path | None,
		typer.Option(
			"--calibration-dataset-path",
			help="Path with assets for the calibration dataset (when required).",
		),
	] = None,
	# verification
	verify: Annotated[
		bool,
		typer.Option(
			help="Run inference before/after conversion and compare.",
		),
	] = False,
	verify_dataset: Annotated[
		str | None,
		typer.Option(
			"--verify-dataset",
			help="Dataset name used for verification. Automatically sets `--verify` to True.",
			show_default="None; Random if `--verify` is set",
		),
	] = None,
	verify_dataset_path: Annotated[
		Path | None,
		typer.Option(
			"--verify-dataset-path",
			help="Path with assets for the verification dataset.",
		),
	] = None,
	# common
	artifacts: ArtifactsDirType = Path("build"),
	use_remote: SeparateEnvironmentType = False,
	repository_path: RepositoryPathType = Path("."),
	output: OutputType = Path("."),
	verbose: VerboseType = [],
):
	"""
	Convert an existing model to a different format. Target format can be quantized.
	"""
	assert quantization not in NEED_CALIBRATION or calibration_dataset is not None

	artifacts.mkdir(parents=True, exist_ok=True)
	output.mkdir(parents=True, exist_ok=True)

	metadata_list = asyncio.run(get_local_metadata(repository_path))
	metadata = next(x for x in metadata_list if x.name == model_name)

	if calibration_dataset is not None:
		parsed_calibration_dataset = (
			_get_dataset(
				calibration_dataset, calibration_dataset_path, split=DatasetSplit.val
			)
			.shuffle(42)
			.take(1000)
		)
	else:
		parsed_calibration_dataset = None

	if verify or verify_dataset is not None:
		parsed_verify_dataset = (
			_get_dataset(
				verify_dataset or "Random", verify_dataset_path, split=DatasetSplit.val
			)
			.shuffle(42)
			.take(1000)
		)
	else:
		parsed_verify_dataset = None

	try:
		results = asyncio.run(
			Convert(
				metadata.format,
				ExactModelFormat(model_format, quantization, hw_acceleration),
				calibration_dataset=parsed_calibration_dataset,
				verify_dataset=parsed_verify_dataset,
				override=override,
				use_remote=use_remote,
				error_immediately=True,
				metadata=metadata,
				output=output,
				artifacts=artifacts,
				default_model=weights_path,
			)()
		)
		if results is not None:
			print(results)
	except RemoteError as e:
		handle_remote_error(e)


@app.command()
def inference(
	weights_path: WeightsPathType,
	model_name: ModelNameTypeOptional = None,
	dataset_name: Annotated[
		str,
		typer.Argument(help="Dataset used for inference."),
	] = "Random",
	dataset_path: Annotated[
		Path | None,
		typer.Option(
			help="Optional path to the dataset assets (required for non-random datasets)."
		),
	] = None,
	take: Annotated[
		int,
		typer.Option(
			help="Number of samples to run through the model during inference.",
		),
	] = 10,
	split: Annotated[
		DatasetSplit, typer.Option(help="Dataset split to use.")
	] = DatasetSplit.val,
	# common
	artifacts: ArtifactsDirType = Path("build"),
	output: OutputTypeOptional = None,
	use_remote: SeparateEnvironmentType = False,
	repository_path: RepositoryPathType = Path("."),
	verbose: VerboseType = [],
):
	"""
	Run inference on a model using a selected dataset and print results.
	"""
	metadata_list = asyncio.run(get_local_metadata(repository_path))
	if model_name:
		metadata = next(x for x in metadata_list if x.name == model_name)
	else:
		assert weights_path.suffix == ".json", (
			f"Model name is required unless using a conversion result (weights_path: {weights_path})."
		)
		metadata = metadata_list[0]

	dataset = _get_dataset(dataset_name, dataset_path, split).take(take)
	try:
		results = asyncio.run(
			run_inference(
				weights_path,
				metadata,
				dataset,
				artifacts,
				output,
				{},
				use_remote,
			)
		)

		pprint(results)

	except RemoteError as e:
		handle_remote_error(e)


@app.callback(invoke_without_command=True)
def main(
	ctx: typer.Context,
	version: Annotated[
		bool, typer.Option("--version", "-v", help="Print version")
	] = False,
):
	if version:
		import importlib.metadata

		version_string_of_foo = importlib.metadata.version("ai-parade")
		print(version_string_of_foo)
		raise typer.Exit()
	if ctx.invoked_subcommand is None:
		pass


if __name__ == "__main__":
	app()
