from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List

import datasets
from datasets.exceptions import DatasetGenerationError
import tempfile


@dataclass
class LocalSplitSpec:
    """Specification for a single local split to be published to HF."""

    name: str
    format: str  # "jsonl", "parquet", "csv"
    paths: List[Path]


def _builder_and_files(split: LocalSplitSpec) -> tuple[str, Dict[str, List[str]]]:
    fmt = split.format.lower()
    if fmt == "jsonl":
        builder = "json"
    elif fmt == "parquet":
        builder = "parquet"
    elif fmt == "csv":
        builder = "csv"
    else:
        raise ValueError(f"Unsupported format for publish-hf: {split.format!r}")

    data_files = {split.name: [str(path) for path in split.paths]}
    return builder, data_files


def publish_hf_dataset(
    dataset_id: str,
    split: LocalSplitSpec,
    *,
    private: bool = True,
    token: str | None = None,
    max_shard_size: str | None = None,
    image_column: str | None = None,
) -> None:
    """
    Build a Hugging Face dataset from local files and push it to the Hub.

    Parameters
    ----------
    dataset_id:
        Target dataset id on the Hugging Face Hub, e.g. "owner/name".
    split:
        Description of the local split to publish (format + paths).
    private:
        Whether to create/update the repo as private. Defaults to True.
    token:
        Optional HF token. When omitted, relies on the user's HF CLI login
        or environment configuration.
    """
    fmt = split.format.lower()

    if fmt == "jsonl":
        # For JSONL, prefer the standard datasets path but fall back to a
        # JSONL -> Parquet conversion when schema casting fails. This avoids
        # issues with complex nested structs (like large head_hist maps) while
        # keeping memory usage reasonable.
        builder, data_files = _builder_and_files(split)
        try:
            ds_dict = datasets.load_dataset(builder, data_files=data_files)
        except (DatasetGenerationError, TypeError, ValueError):
            try:
                import pyarrow.json as paj  # type: ignore[import]
                import pyarrow.parquet as pq  # type: ignore[import]
            except Exception as exc:  # pragma: no cover - import failure path
                raise RuntimeError(
                    "Failed to build dataset from JSONL using datasets.load_dataset "
                    "and pyarrow is not available for the fallback JSONL->Parquet "
                    "conversion path."
                ) from exc

            with tempfile.TemporaryDirectory(prefix="hfschema_jsonl_to_parquet_") as tmpdir:
                parquet_paths: list[str] = []
                for idx, path in enumerate(split.paths):
                    table = paj.read_json(str(path))
                    out = Path(tmpdir) / f"part-{idx:05d}.parquet"
                    pq.write_table(table, out)
                    parquet_paths.append(str(out))
                ds_dict = datasets.load_dataset(
                    "parquet", data_files={split.name: parquet_paths}
                )
    else:
        builder, data_files = _builder_and_files(split)
        ds_dict = datasets.load_dataset(builder, data_files=data_files)

    # Optional: convert an image path column into an Image feature so that
    # actual image assets are uploaded to HF instead of just path strings.
    if image_column and split.format.lower() == "parquet":
        try:
            from datasets import Image
        except Exception:
            raise ValueError(
                "Image feature support requires the vision extras for `datasets`. "
                "Install with `pip install 'datasets[vision]'` or disable image upload."
            )

        for split_name in list(ds_dict.keys()):
            ds_split = ds_dict[split_name]

            def _add_image(example, col=image_column):
                return {"image": example[col]}

            ds_split = ds_split.map(_add_image)
            ds_split = ds_split.cast_column("image", Image())
            # Drop the original path column so HF does not retain local
            # filesystem paths (which can leak usernames or directory layout).
            if image_column in ds_split.column_names:
                ds_split = ds_split.remove_columns([image_column])
            ds_dict[split_name] = ds_split

    # Let datasets/huggingface_hub handle repo creation and versioning.
    push_kwargs: dict = {"private": private, "token": token}
    if max_shard_size is not None:
        push_kwargs["max_shard_size"] = max_shard_size
    ds_dict.push_to_hub(dataset_id, **push_kwargs)
