import logging
import os
from pathlib import Path
from typing import Any, Callable, override

import xmltodict

from ai_parade._toolkit.data.output import ModelOutput

from .Dataset import DatasetSplit
from .FileDataset import FileDataset

logger = logging.getLogger(__name__)


class PascalVOC(FileDataset):
	def __init__(
		self,
		dataset_path: Path,
		split: DatasetSplit,
		annotation_mapping: Callable[[Any], ModelOutput],
		annotations_dir: str = "Annotations",
		data_dir: str = "Data",
	):
		super().__init__(annotation_mapping)
		self._dataset_path = dataset_path
		match split:
			case DatasetSplit.train:
				split_dir = "train"
			case DatasetSplit.val:
				split_dir = "val"
			case DatasetSplit.test:
				split_dir = "test"

		# handle cases like: ILSVRC/Data/VID when there are no other datasets (CLS, DET, VID, etc.)
		data_content = list((dataset_path / data_dir).iterdir())
		if len(data_content) == 1:
			self._data_path = data_content[0] / split_dir
		else:
			self._data_path = dataset_path / data_dir / split_dir

		annotations_content = list((dataset_path / annotations_dir).iterdir())
		if len(annotations_content) == 1:
			self._annotations_path = annotations_content[0] / split_dir
		else:
			self._annotations_path = dataset_path / annotations_dir / split_dir

		if not self._data_path.exists():
			raise ValueError(f"Data directory {self._data_path} does not exist")
		if not self._annotations_path.exists():
			raise ValueError(
				f"Annotations directory {self._annotations_path} does not exist"
			)

	@override
	def walk(self):
		"""Walk the dataset directory tree.

		Returns:
			 iterator of tuples of (path to the annotation file and path to the data file)
		"""
		""" Example paths:
		- Data/VID/test/ILSVRC2015_test_00000000/000000.JPEG
		- Data/VID/train/ILSVRC2015_VID_train_0000/ILSVRC2015_train_00000000/000000.JPEG
		- Annotations/VID/train/ILSVRC2015_VID_train_0000/ILSVRC2015_train_00000000/000000.xml
		"""
		data_path_len = len(self._data_path.parts)
		for path, dirs, files in os.walk(self._data_path):
			# this makes the walking of the directories in sorted order
			# https://stackoverflow.com/a/18282602/5739006
			dirs.sort()
			for file in sorted(files):
				data_path = Path(path) / file
				# this is basically data_path.relative_to(self._data_path)
				# but it runs much faster
				relative_data_path = Path(*data_path.parts[data_path_len:])
				annotation_path = (
					self._annotations_path / relative_data_path.with_suffix(".xml")
				)
				if annotation_path.exists():
					with open(annotation_path) as fd:
						annotation = xmltodict.parse(fd.read())
				else:
					logger.warning(f"Annotation file {annotation_path} does not exist")
					annotation = {}
				yield data_path, annotation
