#!/usr/bin/env python

"""████
██    ██    Datature
██  ██    Powering Breakthrough AI
██

@File    :   loader.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   HuggingFace dataset loader for Vi platform integration.
"""

from __future__ import annotations

import functools
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from collections.abc import Generator

from datasets import (
    Dataset,
    DatasetDict,
    IterableDataset,
    IterableDatasetDict,
    concatenate_datasets,
    interleave_datasets,
    load_dataset,
)
from huggingface_hub import HfApi, HfFileSystem
from PIL import Image, UnidentifiedImageError
from rich.progress import BarColumn, TimeElapsedColumn
from vi import Client
from vi.api.resources.datasets.assets.consts import (
    MAX_PARALLEL_BATCHES,
    MAX_UPLOAD_BATCH_SIZE,
)
from vi.api.resources.datasets.assets.responses import FailureMode
from vi.api.resources.datasets.assets.results import AssetUploadResult
from vi.client.errors import ViError, ViValidationError
from vi.dataset.loaders.hf import consts
from vi.dataset.loaders.hf.batch_processor import (
    HFBatchProcessor,
    extract_image_from_example,
    generate_batches,
)
from vi.dataset.loaders.hf.uploader import ImageFormat
from vi.utils.graceful_exit import graceful_exit
from vi.utils.progress import ViProgress

logger = logging.getLogger(__name__)


class HFDatasetError(ViError):
    """Error raised when HuggingFace dataset operations fail."""


class HFDatasetLoader:
    """Loader for HuggingFace datasets with Vi platform integration.

    This class provides functionality to:
    - Load datasets from HuggingFace Hub
    - Stream images from datasets
    - Upload images to Datature Vi platform
    - Handle various dataset formats and splits
    - Fallback to raw file scraping if metadata is corrupted
    """

    def __init__(
        self,
        dataset_name: str,
        split: str | None = None,
        streaming: bool = True,
        cache_dir: str | Path | None = None,
        trust_remote_code: bool = False,
        **kwargs: Any,
    ):
        """Initialize HuggingFace dataset loader."""
        self.dataset_name = dataset_name
        self.split = split
        self.streaming = streaming
        self.cache_dir = Path(cache_dir) if cache_dir else None
        self.trust_remote_code = trust_remote_code
        self.load_dataset_kwargs = kwargs
        self.dataset: Dataset | DatasetDict | IterableDataset | IterableDatasetDict

        logger.info(
            f"Loading HuggingFace dataset: {dataset_name}"
            + (f" (split: {split})" if split else "")
        )

        try:
            # Attempt 1: Standard load
            self.dataset = load_dataset(
                dataset_name,
                split=split,
                streaming=streaming,
                cache_dir=str(self.cache_dir) if self.cache_dir else None,
                trust_remote_code=trust_remote_code,
                **self.load_dataset_kwargs,
            )
        except ValueError as e:
            # Attempt 2: Handle mixed formats (e.g., imagefolder + webdataset)
            if "Couldn't infer the same data file format" in str(e) and split is None:
                logger.warning(
                    f"Mixed formats detected in '{dataset_name}'. "
                    "Attempting to discover and load splits individually..."
                )
                try:
                    self.dataset = self._load_individual_splits(error_msg=str(e))
                except HFDatasetError as err:
                    logger.warning(f"Individual split loading failed: {err}")
                    logger.info("Falling back to raw image scraping...")
                    self.dataset = self._scrape_dataset()
            else:
                # If it's a different ValueError, try scraping as last resort
                logger.warning(
                    f"Standard loading failed: {e}. Falling back to scraping..."
                )
                self.dataset = self._scrape_dataset()

        except Exception as e:
            # Catch-all for other errors (e.g., XML parsing, 404s on config)
            logger.warning(f"Standard loading failed: {e}. Falling back to scraping...")
            self.dataset = self._scrape_dataset()

        logger.info(f"Successfully loaded dataset: {dataset_name}")

    @staticmethod
    def _fetch_single_image(file_path: str, fs: HfFileSystem) -> dict[str, Any] | None:
        """Worker function to fetch a single image from filesystem.

        Args:
            file_path: The path to the file in the HF repository.
            fs: The initialized HfFileSystem object.

        Returns:
            Dictionary containing image object and metadata, or None if failed.

        """
        try:
            with fs.open(file_path, "rb") as f:
                image = Image.open(f)
                image.load()  # Force load data while file is open

                return {
                    "image": image,
                    "file_path": file_path,
                    "filename": file_path.split("/")[-1],
                }
        except (UnidentifiedImageError, OSError, Exception) as e:
            logger.warning(f"Failed to read image {file_path}: {e}")
            return None

    @staticmethod
    def _scraped_image_generator(
        files: list[str], token: str | None
    ) -> Generator[dict[str, Any]]:
        """Fetch images in parallel as a generator.

        Args:
            files: List of file paths to fetch.
            token: HF API token for authentication.

        Yields:
            Dictionaries containing image data.

        """
        # Re-initialize filesystem here to avoid pickling issues with generators
        fs = HfFileSystem(token=token)

        # Create a partial function to pass the filesystem object to the worker
        fetch_fn = functools.partial(HFDatasetLoader._fetch_single_image, fs=fs)

        with ThreadPoolExecutor(
            max_workers=consts.MAX_PARALLEL_PROCESSING_WORKERS
        ) as executor:
            # executor.map yields results in the order of the input iterable
            for result in executor.map(fetch_fn, files):
                if result is not None:
                    yield result

    def _scrape_dataset(self) -> IterableDatasetDict:
        """Manually scrape image files from the repository via Filesystem API.

        This bypasses the `datasets` library's config inference and directly
        streams image files found in the repository. This is useful for
        malformed datasets or those with mixed/conflicting structures.
        """
        token = self.load_dataset_kwargs.get("token")
        fs = HfFileSystem(token=token)

        # Glob pattern: "datasets/username/repo_name/**/*"
        base_path = f"datasets/{self.dataset_name}"

        logger.info(f"Scraping image files from {base_path}...")

        try:
            # We use a broad glob and filter in python to avoid multiple API calls
            all_files = fs.glob(f"{base_path}/**/*", detail=False)

            # Filter for images
            files = [
                f
                for f in all_files
                if f.split(".")[-1].lower() in consts.IMAGE_EXTENSIONS
            ]

            if not files:
                raise HFDatasetError(
                    f"No image files found in {self.dataset_name} via scraping."
                )

            logger.info(
                f"Found {len(files)} image files via scraping. "
                "Starting parallel download..."
            )

        except Exception as e:
            raise HFDatasetError(f"Failed to scrape dataset files: {e}") from e

        # Create the IterableDataset using the static generator method
        # We pass arguments via gen_kwargs to ensure clean serialization
        dataset = IterableDataset.from_generator(
            self._scraped_image_generator,
            gen_kwargs={"files": files, "token": token},
        )

        # Wrap in dict to mimic standard structure (using 'scraped' as split name)
        return IterableDatasetDict({"scraped": dataset})

    def _load_individual_splits(
        self, error_msg: str = ""
    ) -> DatasetDict | IterableDatasetDict:
        """Load available splits individually and combine them."""
        potential_splits = []

        # Strategy 1: Extract from error message
        if error_msg:
            found = re.findall(r"NamedSplit\('([^']+)'\)", error_msg)
            if found:
                potential_splits.extend(found)

        # Strategy 2: Fetch from Hub
        if not potential_splits:
            hub_splits = self._get_available_splits()
            potential_splits.extend(hub_splits)

        # Strategy 3: Common defaults
        if not potential_splits:
            common_splits = ["train", "validation", "test", "val", "dev"]
            potential_splits.extend(common_splits)

        unique_splits = list(dict.fromkeys(potential_splits))
        loaded_splits = {}

        for split_name in unique_splits:
            try:
                loaded_splits[split_name] = load_dataset(
                    self.dataset_name,
                    split=split_name,
                    streaming=self.streaming,
                    cache_dir=str(self.cache_dir) if self.cache_dir else None,
                    trust_remote_code=self.trust_remote_code,
                    **self.load_dataset_kwargs,
                )
                logger.info(f"Successfully loaded split: {split_name}")
            except Exception as e:
                logger.debug(f"Skipping split '{split_name}' (failed to load): {e}")

        if not loaded_splits:
            raise HFDatasetError("Failed to load any splits via individual loading.")

        if self.streaming:
            return IterableDatasetDict(loaded_splits)
        return DatasetDict(loaded_splits)

    def _get_available_splits(self) -> list[str]:
        """Fetch available split names from the Hugging Face Hub."""
        try:
            api = HfApi()
            info = api.dataset_info(repo_id=self.dataset_name)

            if not info.card_data:
                return []

            try:
                dataset_info = info.card_data["dataset_info"]
            except KeyError:
                return []

            if isinstance(dataset_info, list):
                dataset_info = dataset_info[0] if dataset_info else {}

            splits_data = dataset_info.get("splits", [])
            return [s["name"] for s in splits_data if "name" in s]

        except Exception as e:
            logger.warning(f"Error fetching split info from Hub: {e}")
            return []

    def _get_split_dataset(
        self, split_name: str | None = None
    ) -> Dataset | IterableDataset:
        """Get dataset for a specific split.

        Args:
            split_name: Name of the split to retrieve.

        Returns:
            Dataset or IterableDataset for the specified split.

        """
        if isinstance(self.dataset, (dict, DatasetDict, IterableDatasetDict)):
            if split_name is None:
                # If only one split exists (e.g. 'scraped'), return it automatically
                keys = list(self.dataset.keys())
                if len(keys) == 1:
                    return self.dataset[keys[0]]

                msg = (
                    f"Dataset has multiple splits: {keys}. "
                    "Please specify split_name parameter."
                )
                raise HFDatasetError(msg)

            if split_name not in self.dataset:
                available_splits = list(self.dataset.keys())
                msg = (
                    f"Split '{split_name}' not found in dataset. "
                    f"Available splits: {available_splits}"
                )
                raise HFDatasetError(msg)

            return self.dataset[split_name]

        if split_name is not None and split_name != self.split:
            msg = (
                f"Dataset was loaded with split '{self.split}', "
                f"but '{split_name}' was requested."
            )
            raise HFDatasetError(msg)
        return self.dataset

    def stream_images(
        self,
        image_column: str = "image",
        split_name: str | None = None,
        limit: int | None = None,
        batch_size: int = consts.MAX_PARALLEL_PROCESSING_WORKERS,
        allowed_formats: list[str | ImageFormat] | None = None,
    ) -> list[tuple[Image.Image, dict[str, Any]]]:
        """Stream images from the dataset, optionally filtering by storage format."""
        dataset = self._get_split_dataset(split_name)
        images: list[tuple[Image.Image, dict[str, Any]]] = []

        valid_formats = None
        if allowed_formats:
            valid_formats = {
                f.value.upper() if isinstance(f, ImageFormat) else f.upper()
                for f in allowed_formats
            }
            logger.info(f"Streaming images restricted to formats: {valid_formats}")
        else:
            logger.info(f"Streaming images from column '{image_column}' (all formats)")

        count = 0
        for batch in dataset.iter(batch_size=batch_size):
            if limit is not None and count >= limit:
                break

            batch_size_actual = len(batch[next(iter(batch))])
            for i in range(batch_size_actual):
                if limit is not None and count >= limit:
                    break

                example = {key: batch[key][i] for key in batch}
                image = extract_image_from_example(example, image_column)

                if image is not None:
                    if valid_formats:
                        img_fmt = image.format.upper() if image.format else None
                        if img_fmt and img_fmt not in valid_formats:
                            logger.debug(f"Skipping image with format {img_fmt}")
                            continue

                    images.append((image, example))
                    count += 1

        logger.info(f"Streamed {count} images")
        return images

    def upload(
        self,
        client: Client,
        dataset_id: str,
        image_column: str = "image",
        split_name: str | None = None,
        limit: int | None = None,
        batch_size: int = MAX_UPLOAD_BATCH_SIZE,
        image_format: ImageFormat = ImageFormat.PNG,
        failure_mode: FailureMode = FailureMode.FAIL_AFTER_ALL,
        on_asset_overwritten: str = "RemoveLinkedResource",
        wait_until_done: bool = False,
        show_progress: bool = True,
        parallel: bool = True,
    ) -> AssetUploadResult:
        """Upload images from HuggingFace dataset."""
        if not dataset_id:
            raise ViValidationError("dataset_id is required")

        # Determine which dataset/splits to use
        dataset: Dataset | IterableDataset
        if split_name:
            dataset = self._get_split_dataset(split_name)
        elif isinstance(self.dataset, (dict, DatasetDict, IterableDatasetDict)):
            available_splits = list(self.dataset.keys())

            # Optimization: If we only have one split (e.g. from scraping), just use it
            if len(available_splits) == 1:
                dataset = self.dataset[available_splits[0]]
            else:
                logger.info(
                    f"No split specified. Combining and uploading all splits: {available_splits}"
                )
                try:
                    datasets_list = [self.dataset[s] for s in available_splits]

                    if self.streaming:
                        dataset = interleave_datasets(
                            datasets_list, stop_on_exhaustion=True
                        )
                    else:
                        dataset = concatenate_datasets(datasets_list)

                except Exception as e:
                    raise HFDatasetError(
                        f"Failed to combine dataset splits: {e}"
                    ) from e
        else:
            dataset = self.dataset

        logger.info(
            f"Starting batch upload (batch_size: {batch_size}, limit: {limit or 'all'}, parallel: {parallel})"
        )

        with graceful_exit("Upload cancelled by user") as handler:
            try:
                config = {
                    "dataset_id": dataset_id,
                    "failure_mode": failure_mode,
                    "on_asset_overwritten": on_asset_overwritten,
                    "wait_until_done": wait_until_done,
                }

                dataset_size = self._get_dataset_size()
                estimated_batches = self._estimate_batch_count(
                    dataset_size, limit, batch_size
                )
                total_images_estimate = limit if limit else dataset_size

                progress = None
                if show_progress:
                    progress = self._create_progress_display()

                batch_generator = generate_batches(
                    dataset=dataset,
                    image_column=image_column,
                    batch_size=batch_size,
                    limit=limit,
                    handler=handler,
                    image_format=image_format,
                    progress=progress,
                    total_images=total_images_estimate,
                )

                processor = HFBatchProcessor(client, config, progress, handler)
                results = []

                if show_progress and progress:
                    with progress:
                        if parallel and estimated_batches > 1:
                            results = processor.process_batches_parallel(
                                batch_generator=batch_generator,
                                max_parallel=MAX_PARALLEL_BATCHES,
                            )
                        else:
                            results = processor.process_batches_sequential(
                                batch_generator=batch_generator,
                            )
                else:
                    if parallel and estimated_batches > 1:
                        results = processor.process_batches_parallel(
                            batch_generator=batch_generator,
                            max_parallel=MAX_PARALLEL_BATCHES,
                        )
                    else:
                        results = processor.process_batches_sequential(
                            batch_generator=batch_generator,
                        )

                if not results:
                    msg = (
                        "No images could be uploaded. Possible causes:\n"
                        "  1. No images found in the specified dataset/split\n"
                        "  2. Image column not found or empty\n"
                        "  3. All images failed to save to disk\n"
                        "  4. Upload was cancelled"
                    )
                    raise HFDatasetError(msg)

                all_sessions = []
                total_images = 0
                for result in results:
                    all_sessions.extend(result.sessions)
                    total_images += result.total_succeeded + result.total_failed

                logger.info(
                    f"Upload complete: {total_images} images processed in {len(results)} batch(es)"
                )

                return AssetUploadResult(sessions=all_sessions)

            except Exception as e:
                raise HFDatasetError(f"Failed to upload images: {e}") from e

    def _get_dataset_size(self) -> int | None:
        """Get the actual size of a HuggingFace dataset."""
        if not self.streaming:
            try:
                if isinstance(self.dataset, (DatasetDict, IterableDatasetDict)):
                    if self.split and self.split in self.dataset:
                        return len(self.dataset[self.split])
                    return sum(len(d) for d in self.dataset.values())
                if hasattr(self.dataset, "__len__"):
                    return len(self.dataset)
            except (TypeError, AttributeError):
                pass

        try:
            api = HfApi()
            info = api.dataset_info(repo_id=self.dataset_name)

            if not info.card_data:
                return None

            try:
                dataset_info = info.card_data["dataset_info"]
            except KeyError:
                return None

            if isinstance(dataset_info, list):
                if not dataset_info:
                    return None
                dataset_info = dataset_info[0]

            splits = dataset_info.get("splits", [])

            if self.split:
                for split in splits:
                    if split.get("name") == self.split:
                        return split.get("num_examples")
                return None

            total_examples = sum(split.get("num_examples", 0) for split in splits)
            return total_examples

        except Exception as e:
            logger.debug(f"Could not determine dataset size from HF Hub: {e}")

        return None

    def _estimate_batch_count(
        self, dataset_size: int | None, limit: int | None, batch_size: int
    ) -> int:
        """Estimate the number of batches for progress tracking."""
        effective_size = limit if limit is not None else dataset_size
        if effective_size is not None:
            return max(1, (effective_size + batch_size - 1) // batch_size)
        return 5

    def _create_progress_display(self) -> ViProgress:
        """Create the progress bar display for upload tracking."""
        return ViProgress(
            "[progress.description]{task.description}",
            BarColumn(),
            "[progress.percentage]{task.percentage:>3.2f}%",
            TimeElapsedColumn(),
            transient=False,
        )

    def get_splits(self) -> list[str]:
        """Get available dataset splits."""
        if isinstance(self.dataset, (dict, DatasetDict, IterableDatasetDict)):
            return list(self.dataset.keys())
        return [self.split] if self.split else ["default"]

    def get_info(self) -> dict[str, Any]:
        """Get dataset information."""
        info: dict[str, Any] = {
            "dataset_name": self.dataset_name,
            "splits": self.get_splits(),
            "streaming": self.streaming,
        }
        if hasattr(self.dataset, "info"):
            info["features"] = str(self.dataset.info.features)
            info["description"] = self.dataset.info.description
            info["homepage"] = self.dataset.info.homepage
        return info

    def __repr__(self) -> str:
        """Return string representation of the loader."""
        return (
            f"HFDatasetLoader("
            f"dataset_name='{self.dataset_name}', "
            f"split={self.split}, "
            f"streaming={self.streaming})"
        )
