from __future__ import annotations

from dataclasses import dataclass
import importlib
import importlib.resources as pkg_resources
import re
from pathlib import Path
import glob
import os

import typer

from .inspect import DatasetInfoError, describe_dataset, summarize_feature
from .publish import LocalSplitSpec, publish_hf_dataset
from .bundles import DATASET_BUNDLES
from .wizard import suggest_plan_for_warp_dataset, PlanKind
from .embeddings import (
    add_space as _add_embedding_space,
    default_embeddings_store,
    link_field as _link_embedding_field,
    load_embeddings_config,
)
from .naming import module_filename, schema_key
from .registry import registry
from .sync import generate_from_cache
from .custom import generate_schema_from_paths
from datasets import load_dataset_builder
from huggingface_hub import HfApi, whoami

app = typer.Typer(
    help=(
        "Manage local Hugging Face dataset schemas (CLI aliases: hfschema / schemas-hf).\n\n"
        "MENTAL MODEL\n"
        "============\n"
        "hfschema sits on top of the `datasets` + HF Hub ecosystem and focuses on one job:\n"
        "turning HF datasets (public or private) or local files into small, typed Python\n"
        "dataclasses you can import in your own code. It never owns the raw data or the\n"
        "training loop; it only manages schemas and simple loader templates.\n\n"
        "There are three main layers:\n"
        "  • HF datasets (public or private): identified by `owner/name[:config]`.\n"
        "  • Local tabular files: JSONL, Parquet, CSV that you want to expose as datasets.\n"
        "  • Generated schema modules: `schemas_hf/datasets/*.py` dataclasses + helpers.\n\n"
        "You typically use hfschema in one of these ways:\n"
        "  1) For an existing HF dataset:\n"
        "       - `hfschema download owner/name` to ensure metadata is cached and\n"
        "         generate schema modules for that dataset only.\n"
        "       - `hfschema list` to see what schemas exist.\n"
        "       - `hfschema instantiate owner/name --dest ./schemas` to copy the\n"
        "         generated dataclass into your own repo.\n"
        "       - (optional) `hfschema init-loader ...` to scaffold a streaming loader.\n\n"
        "  2) For local Parquet/JSONL/CSV you want to publish as an HF dataset first:\n"
        "       - `hfschema publish-hf uwunion/my-dataset --format parquet --path data/*.parquet`\n"
        "         to push a script-free HF dataset (private by default).\n"
        "       - Then `hfschema sync --include uwunion/my-dataset` or\n"
        "         `hfschema download uwunion/my-dataset` to generate schema modules.\n"
        "       - Finally, `instantiate` or `init-loader` as above in your consuming repo.\n\n"
        "  3) For purely local schemas (no HF Hub):\n"
        "       - `hfschema make-local-schema local/thing -f jsonl -p raw/*.jsonl -o ./schemas`.\n"
        "         This inspects a sample of rows and emits a dataclass module without\n"
        "         ever touching HF Hub. You import that module directly.\n\n"
        "WHERE FILES GO\n"
        "==============\n"
        "By default, schema modules are written under the hfschema package itself:\n"
        "  • `sync` and `download` default to: `.../schemas_hf/datasets`.\n"
        "This makes sense when you are *developing hfschema* or curating built-in schemas.\n\n"
        "For project-specific schemas, you can and should override the output directory:\n"
        "  • `hfschema download owner/name --output ./schemas_hf_datasets`\n"
        "  • `hfschema sync --include owner/name --output ./schemas_hf_datasets`\n"
        "Then either import the generated modules from that path or copy them into your\n"
        "own package. `instantiate` is a convenient way to copy a single schema from the\n"
        "built-in `schemas_hf.datasets` package into your repo.\n\n"
        "HIGH-LEVEL COMMAND MAP\n"
        "======================\n"
        "  • Introspection:\n"
        "      - `list`   : show all dataset[:config] schemas currently available.\n"
        "      - `info`   : show HF metadata for a dataset (no schema required).\n\n"
        "  • Schema generation:\n"
        "      - `sync`     : generate/refresh schema modules from HF caches. With\n"
        "                     `--include owner/name[:config]` it only touches those ids.\n"
        "      - `download` : one-step `load_dataset_builder` + filtered `sync` for\n"
        "                     the given dataset ids (pull from HF and write schemas).\n"
        "      - `make-local-schema` : infer a schema from local JSONL (and in future\n"
        "                     other formats) without using HF Hub.\n\n"
        "  • Publishing local data as HF datasets:\n"
        "      - `publish-hf` : turn local JSONL/Parquet/CSV into a script-free HF\n"
        "                       dataset (public or private), which you can then feed\n"
        "                       back through `sync`/`download` like any other HF dataset.\n\n"
        "  • Reusing schemas in your own repos:\n"
        "      - `instantiate` : copy a generated schema module into another directory.\n"
        "      - `init-loader` : scaffold a Python loader that streams multiple\n"
        "                        datasets and uses their schemas for validation.\n\n"
        "  • Embeddings metadata (optional):\n"
        "      - `embeddings add`        : define a named embedding space (model + provider).\n"
        "      - `embeddings list`       : list known embedding spaces.\n"
        "      - `embeddings link-field` : link a dataset field to a named space so other\n"
        "                                   tools can discover \"dataset.field → model\".\n\n"
        "Remember: you only need `sync`/`download` when you want *schema files* or\n"
        "auto-generated loaders. If you just want to load data with `datasets`, you\n"
        "can continue to call `datasets.load_dataset(...)` directly without hfschema."
    )
)


@app.callback()
def main() -> None:
    """Schemas HF command line interface."""


embeddings_app = typer.Typer(help="Manage embedding spaces and schema links.")
app.add_typer(embeddings_app, name="embeddings")


@app.command()
def sync(
    cache_dir: Path = typer.Option(
        Path.home() / ".cache" / "huggingface" / "hub",
        "--cache-dir",
        "-c",
        help="Path to the Hugging Face hub cache.",
    ),
    output: Path = typer.Option(
        Path(__file__).resolve().parent / "datasets",
        "--output",
        "-o",
        help="Directory where schema modules will be written.",
    ),
    include: list[str] = typer.Option(
        None,
        "--include",
        "-i",
        help="Dataset ids to include (optionally with ':config'). Repeat to select multiple.",
    ),
    dry_run: bool = typer.Option(False, "--dry-run", help="Inspect actions without writing files."),
) -> None:
    """
    Materialize schema modules from your local HF caches.

    Examples:
      • Sync everything using defaults:
          `schemas-hf sync`
      • Only generate ag_news + glue/mrpc:
          `schemas-hf sync -i ag_news -i glue:mrpc`
      • Custom hubs (e.g., remote cache mount):
          `schemas-hf sync -c /mnt/hf_hub --output src/my_pkg/datasets`
    """
    only = set(include) if include else None
    # For broad syncs, we default to skipping datasets that already have a
    # generated schema module in the output directory. Use --include to force
    # regeneration for specific ids or override the output path.
    generated = generate_from_cache(
        cache_dir,
        output,
        dry_run=dry_run,
        only=only,
        skip_existing=True,
    )
    typer.echo(f"Generated {len(generated)} schema file{'s' if len(generated) != 1 else ''} at {output}")


@app.command("sync-space")
def sync_space(
    space: str | None = typer.Option(
        None,
        "--space",
        "-s",
        help=(
            "HF username or organization whose datasets should be synced. "
            "Defaults to the current HF user from `huggingface-cli whoami`."
        ),
    ),
    cache_dir: Path = typer.Option(
        Path.home() / ".cache" / "huggingface" / "hub",
        "--cache-dir",
        "-c",
        help="Path to the Hugging Face hub cache.",
    ),
    output: Path = typer.Option(
        Path(__file__).resolve().parent / "datasets",
        "--output",
        "-o",
        help="Directory where schema modules will be written.",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Only download metadata and report actions without writing schema files.",
    ),
) -> None:
    """
    Discover all datasets under a HF user/org and generate schemas for them.

    This command is a convenience wrapper around:
      1. Listing datasets for a given HF space (user or organization).
      2. Calling `datasets.load_dataset_builder` for each to ensure metadata
         is cached locally.
      3. Running a filtered `sync` for just those dataset ids.
    """
    # Resolve the namespace to use.
    if space is None:
        try:
            info = whoami()
            namespace = info.get("name") or info.get("username")
        except Exception as exc:  # pragma: no cover - depends on local HF setup
            typer.echo(
                "Unable to determine default HF user from `huggingface-cli whoami`.\n"
                "Please pass an explicit --space USER_OR_ORG.",
                err=True,
            )
            typer.echo(f"Reason: {exc}", err=True)
            raise typer.Exit(code=1)
    else:
        namespace = space

    api = HfApi()
    try:
        entries = api.list_datasets(author=namespace)
    except Exception as exc:  # pragma: no cover - network/hub dependent
        typer.echo(
            f"Unable to list datasets for space '{namespace}'. "
            "Check your HF token and network.",
            err=True,
        )
        typer.echo(f"Reason: {exc}", err=True)
        raise typer.Exit(code=1)

    dataset_ids = [entry.id for entry in entries]
    if not dataset_ids:
        typer.echo(f"No datasets found for space '{namespace}'.", err=True)
        raise typer.Exit(code=1)

    # Ensure metadata is downloaded for each dataset id, but skip any repos
    # that do not yet expose usable data files (e.g. README-only stubs).
    successful_ids: list[str] = []
    for ds_id in dataset_ids:
        try:
            load_dataset_builder(ds_id)
        except Exception as exc:  # pragma: no cover - depends on remote datasets
            typer.echo(f"Unable to download metadata for '{ds_id}'. Skipping.", err=True)
            typer.echo(f"Reason: {exc}", err=True)
            continue
        successful_ids.append(ds_id)

    if not successful_ids:
        typer.echo(
            f"Unable to download metadata for any dataset under '{namespace}'. "
            "Ensure at least one dataset has data files and you have access.",
            err=True,
        )
        raise typer.Exit(code=1)

    # Generate or refresh schema modules only for this space's datasets.
    try:
        generated = generate_from_cache(
            cache_dir,
            output,
            dry_run=dry_run,
            only=successful_ids,
            skip_existing=True,
        )
    except Exception as exc:  # pragma: no cover - defensive; surfaced to user
        typer.echo(
            f"Unable to create schema for one or more datasets under '{namespace}'.",
            err=True,
        )
        typer.echo(f"Reason: {exc}", err=True)
        raise typer.Exit(code=1)

    typer.echo(
        f"Synced {len(successful_ids)} dataset id(s) for space '{namespace}'. "
        f"Generated {len(generated)} schema file{'s' if len(generated) != 1 else ''} at {output}"
    )


@app.command()
def download(
    datasets: list[str] = typer.Argument(
        ...,
        help="Dataset ids to download and schema-ify, e.g. 'owner/name' or 'owner/name:config'.",
    ),
    cache_dir: Path = typer.Option(
        Path.home() / ".cache" / "huggingface" / "hub",
        "--cache-dir",
        "-c",
        help="Path to the Hugging Face hub cache.",
    ),
    output: Path = typer.Option(
        Path(__file__).resolve().parent / "datasets",
        "--output",
        "-o",
        help="Directory where schema modules will be written.",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Only download metadata and report actions without writing schema files.",
    ),
) -> None:
    """
    Download HF dataset metadata and generate schema modules in one step.

    This is roughly equivalent to calling `datasets.load_dataset_builder`
    for each dataset to ensure metadata is cached, followed by a filtered
    `hfschema sync` for those ids.
    """
    if not datasets:
        typer.echo("At least one dataset id must be provided.", err=True)
        raise typer.Exit(code=1)

    # Expand dataset ids with any known companions (multi-part bundles).
    requested_specs = list(datasets)
    expanded_specs: list[str] = []
    optional_ids: set[str] = set()

    for spec in requested_specs:
        expanded_specs.append(spec)
        base_id, _, _ = spec.partition(":")
        companions = DATASET_BUNDLES.get(base_id, [])
        for comp in companions:
            if comp not in expanded_specs and comp not in requested_specs:
                expanded_specs.append(comp)
                optional_ids.add(comp)
                typer.echo(
                    f"Detected associated dataset '{comp}' for '{spec}'. "
                    "Will attempt to download its metadata as well.",
                )

    datasets = expanded_specs

    # First, ensure metadata is available in the local hub cache.
    successful_specs: list[str] = []
    for spec in datasets:
        ds_id, _, cfg = spec.partition(":")
        cfg = cfg or None
        try:
            kwargs = {"name": cfg} if cfg else {}
            load_dataset_builder(ds_id, **kwargs)
        except Exception as exc:  # pragma: no cover - depends on remote datasets
            if ds_id in optional_ids:
                typer.echo(
                    f"Associated dataset '{spec}' is not available on HF or "
                    "cannot be accessed with the current token. Skipping it.",
                    err=True,
                )
                typer.echo(f"Reason: {exc}", err=True)
                continue
            typer.echo(f"Unable to download metadata for '{spec}'.", err=True)
            typer.echo(f"Reason: {exc}", err=True)
            raise typer.Exit(code=1)
        else:
            successful_specs.append(spec)

    # Then generate schema modules only for the requested ids.
    try:
        generated = generate_from_cache(
            cache_dir,
            output,
            dry_run=dry_run,
            only=successful_specs,
        )
    except Exception as exc:  # pragma: no cover - defensive; surfaced to user
        typer.echo("Unable to create schema for one or more datasets.", err=True)
        typer.echo(f"Reason: {exc}", err=True)
        raise typer.Exit(code=1)

    if not generated and not dry_run:
        typer.echo(
            "No schema files were generated. Check that the dataset exposes feature metadata "
            "and is supported by the installed `datasets` version.",
            err=True,
        )
        raise typer.Exit(code=1)

    typer.echo(
        f"Generated {len(generated)} schema file{'s' if len(generated) != 1 else ''} at {output}"
    )


@app.command()
def instantiate(
    dataset: str = typer.Argument(
        ...,
        help="Dataset id, e.g. 'ag_news', 'owner/name', or 'owner/name:config'.",
    ),
    config: str | None = typer.Option(None, "--config", help="Optional config to target."),
    destination: Path = typer.Option(
        Path.cwd(),
        "--dest",
        "-d",
        help="Directory or file path where the schema file should be written.",
    ),
    force: bool = typer.Option(False, "--force", help="Overwrite if the destination already exists."),
) -> None:
    """
    Copy a generated schema module into another directory.

    Examples:
      • Drop ag_news into the current repo:
          `schemas-hf instantiate ag_news --dest ./schemas`
      • Grab a config-specific variant:
          `schemas-hf instantiate glue --config mrpc --dest api/models`
      • Force-overwrite an existing file:
          `schemas-hf instantiate imdb --config plain_text -d schema.py --force`
      • Use colon syntax instead of --config:
          `schemas-hf instantiate openai/gsm8k:main --dest ./schemas`
    """
    # Support colon syntax in the dataset argument, e.g. "owner/name:config".
    ds_id = dataset
    cfg = config
    if ":" in ds_id and cfg is None:
        base, _, cfg_part = ds_id.partition(":")
        if cfg_part:
            ds_id, cfg = base, cfg_part

    module_name = module_filename(ds_id, cfg)
    resource_path = pkg_resources.files("schemas_hf.datasets") / module_name
    if not resource_path.exists():
        typer.echo(
            f"Schema for {schema_key(ds_id, cfg)} not found. "
            "Run 'schemas-hf sync' first to generate it.",
            err=True,
        )
        raise typer.Exit(code=1)

    contents = resource_path.read_text()
    target_path = _resolve_destination(destination, module_name)
    if target_path.exists() and not force:
        typer.echo(f"{target_path} already exists. Use --force to overwrite.", err=True)
        raise typer.Exit(code=1)

    target_path.parent.mkdir(parents=True, exist_ok=True)
    target_path.write_text(contents)
    typer.echo(f"Wrote schema to {target_path}")


@app.command("list")
def list_schemas() -> None:
    """
    List all dataset[:config] schemas currently generated.

    Example:
      `schemas-hf list`
    """
    registry.discover()
    names = registry.list_names()
    if not names:
        typer.echo("No schemas registered.")
        raise typer.Exit(code=1)
    for name in names:
        typer.echo(name)


@app.command()
def info(
    dataset: str = typer.Argument(..., help="Dataset id, e.g. 'ag_news' or 'owner/name'."),
    config: str | None = typer.Option(None, "--config", help="Optional config to target."),
) -> None:
    """
    Show metadata pulled from Hugging Face for a dataset.

    Accepts fuzzy names:
      • `schemas-hf info ai-arc` → matches `ai2_arc:ARC-Easy`
      • `schemas-hf info lukaemon bbh` → matches `lukaemon/bbh:boolean_expressions`

    Provide --config for explicit configs:
      `schemas-hf info glue --config mrpc`
    """
    dataset_id, resolved_config = _resolve_schema_selection(dataset, config)
    try:
        details = describe_dataset(dataset_id, resolved_config)
    except DatasetInfoError as exc:
        _echo_dataset_error(dataset_id, resolved_config, exc)
        raise typer.Exit(code=1)
    typer.echo(f"Dataset : {details.dataset_id}")
    typer.echo(f"Config  : {details.config_name or 'default'}")
    if details.homepage:
        typer.echo(f"Homepage: {details.homepage}")
    if details.license:
        typer.echo(f"License : {details.license}")
    if details.description:
        typer.echo("\nDescription:")
        typer.echo(f"  {details.description.strip()[:500]}")
    typer.echo("\nColumns:")
    if details.features:
        for name, spec in details.features.items():
            typer.echo(f"  - {name}: {summarize_feature(spec)}")
    else:
        typer.echo("  (No feature metadata available)")
    if details.splits:
        typer.echo("\nSplits:")
        for split, meta in details.splits.items():
            count = meta.get("num_examples")
            typer.echo(f"  - {split}: {count} rows" if count is not None else f"  - {split}")
    if details.builder_configs:
        typer.echo("\nAvailable configs:")
        typer.echo("  " + ", ".join(details.builder_configs))


@app.command("init-loader")
def init_loader(
    name: str = typer.Argument(..., help="Name for the loader script (slug or snake_case)."),
    datasets: list[str] = typer.Option(
        ...,
        "--dataset",
        "-d",
        help="Dataset spec: dataset[:config][@split]. Example: glue:mrpc@validation",
    ),
    output: Path = typer.Option(
        Path("loaders"),
        "--output",
        "-o",
        help="Directory where loader scripts will be written.",
    ),
    streaming: bool = typer.Option(
        True, "--streaming/--no-streaming", help="Use the HF streaming API when loading."
    ),
    force: bool = typer.Option(False, "--force", help="Overwrite an existing loader script."),
) -> None:
    """
    Generate a ready-to-run loader script that streams multiple datasets.

    Example:
      hfschema init-loader text_mix -d ag_news -d glue:mrpc@validation -o loaders/
    """
    if not datasets:
        raise typer.BadParameter("At least one --dataset spec is required.")

    registry.discover()
    available = registry.list_names()
    if not available:
        raise typer.Exit(code=1)

    entries: list[LoaderEntry] = []
    for spec in datasets:
        selection = _parse_dataset_spec(spec)
        query = selection.dataset_id
        if selection.config:
            query = f"{selection.dataset_id}:{selection.config}"
        canonical = _match_schema_name(query, available)
        dataset_id, config_name = _split_schema_key(canonical)
        split = selection.split or "train"
        try:
            entries.append(_build_loader_entry(dataset_id, config_name, split))
        except Exception as exc:  # pragma: no cover - surfaced to user
            typer.echo(f"Failed to prepare dataset '{spec}': {exc}", err=True)
            raise typer.Exit(code=1) from exc

    target_dir = output
    target_dir.mkdir(parents=True, exist_ok=True)
    target_file = target_dir / f"{_slugify_name(name)}.py"
    if target_file.exists() and not force:
        raise typer.BadParameter(f"{target_file} already exists. Use --force to overwrite.")

    contents = _render_loader_template(name, entries, streaming)
    target_file.write_text(contents)
    typer.echo(f"Created loader at {target_file}")


@app.command("make-local-schema")
def make_local_schema(
    dataset: str = typer.Argument(
        ...,
        help="Dataset id for the generated schema, e.g. 'local/fineweb_small'.",
    ),
    format: str = typer.Option(
        ...,
        "--format",
        "-f",
        help="Input format for local files. Currently supported: 'jsonl'.",
    ),
    path: list[Path] = typer.Option(
        ...,
        "--path",
        "-p",
        help=(
            "Path or glob pattern to source files. "
            "Repeat this option to include multiple files or patterns."
        ),
    ),
    output: Path = typer.Option(
        Path(__file__).resolve().parent / "datasets",
        "--output",
        "-o",
        help="Directory where the schema module will be written.",
    ),
    sample_rows: int = typer.Option(
        1000,
        "--sample-rows",
        help="Number of rows to sample for schema inference.",
    ),
) -> None:
    """
    Generate a schema module from local files (e.g. JSONL) without using HF Hub.

    Example:
      hfschema make-local-schema local/fineweb_small \\
        -f jsonl -p data/fineweb_*.jsonl -o src/schemas_hf/datasets
    """
    if not path:
        raise typer.BadParameter("At least one --path/-p must be provided.")

    # Expand glob patterns manually so the command behaves consistently across shells.
    files: list[Path] = []
    for p in path:
        pattern = str(p)
        if any(ch in pattern for ch in "*?[]"):
            if os.path.isabs(pattern):
                # Absolute patterns: use glob.glob directly.
                matches = [Path(m) for m in glob.glob(pattern)]
            else:
                matches = list(Path().glob(pattern))
            files.extend(sorted(matches))
        else:
            files.append(p)

    if not files:
        typer.echo("No files matched the provided --path patterns.", err=True)
        raise typer.Exit(code=1)

    try:
        module_path = generate_schema_from_paths(
            dataset_id=dataset,
            paths=files,
            format=format,
            output_dir=output,
            sample_rows=sample_rows,
        )
    except ValueError as exc:
        typer.echo(str(exc), err=True)
        raise typer.Exit(code=1)

    typer.echo(f"Generated schema for {dataset} at {module_path}")


@app.command("publish-hf")
def publish_hf(
    dataset_id: str = typer.Argument(
        ..., help="Target HF dataset id, e.g. 'owner/name'."
    ),
    format: str = typer.Option(
        ...,
        "--format",
        "-f",
        help="Input format for local files: 'jsonl', 'parquet', or 'csv'.",
    ),
    path: list[Path] = typer.Option(
        ...,
        "--path",
        "-p",
        help="Path or glob pattern to source files. Repeat for multiple patterns.",
    ),
    split: str = typer.Option(
        "train",
        "--split",
        help="Split name for the uploaded data (default: 'train').",
    ),
    private: bool = typer.Option(
        True,
        "--private/--public",
        help="Create or update the dataset as private (default: private).",
    ),
    max_shard_size: str | None = typer.Option(
        "3GB",
        "--max-shard-size",
        help=(
            "Max shard size for push_to_hub, e.g. '3GB' or '1GB'. "
            "Larger shards mean fewer files but bigger individual uploads "
            "(default: 3GB)."
        ),
    ),
    token: str | None = typer.Option(
        None,
        "--token",
        help="Optional HF token; otherwise relies on local HF CLI login / env vars.",
    ),
    images: bool | None = typer.Option(
        None,
        "--images/--no-images",
        help=(
            "When asset-like columns (e.g. 'image_path') are detected in Parquet input, "
            "upload actual image assets to HF using an Image feature instead of only "
            "storing path strings. When omitted and image paths are detected, you'll "
            "be prompted."
        ),
    ),
) -> None:
    """
    Build a Hugging Face dataset from local files and push it to the Hub.
    """
    if not path:
        raise typer.BadParameter("At least one --path/-p must be provided.")

    # Expand glob patterns so behavior is consistent across shells.
    files: list[Path] = []
    for p in path:
        pattern = str(p)
        if any(ch in pattern for ch in "*?[]"):
            if os.path.isabs(pattern):
                matches = [Path(m) for m in glob.glob(pattern)]
            else:
                matches = list(Path().glob(pattern))
            files.extend(sorted(matches))
        else:
            files.append(p)

    if not files:
        typer.echo("No files matched the provided --path patterns.", err=True)
        raise typer.Exit(code=1)

    # Inspect the Parquet schema (if applicable) so we can warn about asset-like
    # columns and optionally convert image paths into an Image feature.
    image_cols = _inspect_asset_columns(format, files)

    image_column: str | None = None
    if image_cols:
        # If the user explicitly requested images, honor that. Otherwise, prompt.
        if images is True:
            image_column = image_cols[0]
        elif images is None:
            if typer.confirm(
                f"Detected image path column(s) {', '.join(image_cols)}. "
                "Upload images to HF as an Image feature (this may be large)?"
            ):
                image_column = image_cols[0]
        # images is False → keep image_column as None and just upload paths.

    spec = LocalSplitSpec(name=split, format=format, paths=files)
    try:
        publish_hf_dataset(
            dataset_id,
            spec,
            private=private,
            token=token,
            max_shard_size=max_shard_size,
            image_column=image_column,
        )
    except ValueError as exc:
        typer.echo(str(exc), err=True)
        raise typer.Exit(code=1)

    typer.echo(
        f"Pushed HF dataset '{dataset_id}' "
        f"(split='{split}', format='{format}')"
    )

    # Best-effort: immediately generate or refresh the local schema module for
    # this dataset so that `hfschema list` sees it without a separate sync.
    try:
        # Ensure metadata is available locally.
        base_id, _, cfg = dataset_id.partition(":")
        cfg = cfg or None
        kwargs = {"name": cfg} if cfg else {}
        load_dataset_builder(base_id, **kwargs)
        # Generate schema for just this dataset id.
        generate_from_cache(
            cache_dir=Path.home() / ".cache" / "huggingface" / "hub",
            output=Path(__file__).resolve().parent / "datasets",
            dry_run=False,
            only=[dataset_id],
            skip_existing=True,
        )
    except Exception:
        # Do not treat schema generation issues as fatal for publishing; users
        # can always run `hfschema download` or `hfschema sync` manually.
        pass


def _inspect_asset_columns(fmt: str, files: list[Path]) -> list[str]:
    fmt = fmt.lower()
    if fmt != "parquet":
        return []
    try:
        import pyarrow.parquet as pq  # type: ignore[import]
    except Exception:
        return []

    try:
        pf = pq.ParquetFile(str(files[0]))
        schema = pf.schema_arrow
    except Exception:
        return []

    col_names = [f.name for f in schema]
    lower = {name.lower(): name for name in col_names}
    asset_keys = []
    for key in ("image_path", "audio_path", "file_path", "rel_path"):
        if key in lower:
            asset_keys.append(lower[key])
    if asset_keys:
        typer.echo(
            "Warning: detected asset-like columns "
            f"{', '.join(asset_keys)} in the source Parquet. "
            "Only the table (including these path strings) will be uploaded "
            "to HF; the underlying files they point to are *not* included.\n"
            "If you want a fully self-contained image/audio dataset on HF, "
            "you will need a dedicated pipeline that uses `datasets.Image` or "
            "similar feature types to upload the raw assets.",
            err=True,
        )
    # Only image-related columns can be upgraded to Image features.
    image_cols = []
    for name in col_names:
        ln = name.lower()
        if "image_path" in ln or (ln.startswith("image") and "path" in ln):
            image_cols.append(name)
    return image_cols


@app.command("wizard")
def wizard(
    warp_dataset: str = typer.Option(
        None,
        "--warp",
        "-w",
        help=(
            "Optional warp dataset name (e.g. 'vision/coco-train' or "
            "'warpdata://neuro/dmt-brains'). When provided, the wizard "
            "prints a suggested HF migration plan based on known patterns."
        ),
    ),
) -> None:
    """
    Interactive helper to plan HF migrations for complex datasets.

    When called without options, this runs a simple text-based wizard
    that:
      1. Asks you to choose a source type (HF / warp / local).
      2. For warp datasets, asks for the dataset name and then prints a
         suggested migration pattern + HF ids.

    Example:
      hfschema wizard                      # interactive
      hfschema wizard --warp neuro/dmt-brains
      hfschema wizard --warp vision/coco-train
    """
    # Interactive mode when no explicit warp dataset is given.
    if warp_dataset is None:
        typer.echo("hfschema wizard")
        typer.echo("================")
        typer.echo("This helper suggests HF migration patterns for complex datasets.\n")
        typer.echo("Choose a source type:")
        typer.echo("  [1] Existing HF dataset (owner/name[:config])")
        typer.echo("  [2] warp dataset (e.g. 'vision/coco-train')")
        typer.echo("  [3] Local files (Parquet/JSONL/CSV)")
        choice = typer.prompt("Enter choice", default="2").strip()

        if choice == "1":
            typer.echo(
                "\nHF dataset flow is not fully wired into the wizard yet.\n"
                "You can use:\n"
                "  - hfschema download owner/name[:config]  # generate schemas\n"
                "to work directly with existing Hub datasets.\n"
            )
            raise typer.Exit(code=0)

        if choice == "3":
            typer.echo(
                "\nLocal files flow is not fully wired into the wizard yet.\n"
                "You can use:\n"
                "  - hfschema publish-hf uwunion/my-dataset -f parquet -p path/*.parquet\n"
                "  - hfschema make-local-schema local/name -f jsonl -p path/*.jsonl\n"
                "for now.\n"
            )
            raise typer.Exit(code=0)

        # Default to warp dataset path.
        warp_dataset = typer.prompt(
            "\nEnter warp dataset name (e.g. 'vision/coco-train' or 'neuro/dmt-brains')"
        )

    plan = suggest_plan_for_warp_dataset(warp_dataset)
    typer.echo(f"Warp dataset : {plan.warp_name}")
    typer.echo(f"Pattern      : {plan.pattern.value}")
    typer.echo("HF dataset(s):")
    for hf_id in plan.hf_ids:
        typer.echo(f"  - {hf_id}")
    if plan.notes:
        typer.echo("\nNotes:")
        typer.echo(f"  {plan.notes}")

    # Offer to show concrete publish commands.
    if typer.confirm("\nShow suggested commands for this plan?", default=True):
        _print_wizard_commands(plan)


def _print_wizard_commands(plan) -> None:
    """
    Print suggested CLI commands for a wizard plan.

    This keeps logic here rather than in schemas_hf.wizard so that we
    can freely use typer for prompting without affecting tests.
    """
    if plan.pattern == PlanKind.PLAIN_TABLE:
        # Try to auto-detect Parquet paths from the warp registry/cache.
        parquet_paths = _resolve_warp_parquet_paths(plan.warp_name)
        local_path: str
        if parquet_paths:
            if len(parquet_paths) == 1:
                auto = parquet_paths[0]
                typer.echo(f"Detected Parquet file: {auto}")
                if typer.confirm("Use this path?", default=True):
                    local_path = str(auto)
                else:
                    local_path = typer.prompt(
                        "Enter local Parquet path for the main table "
                        "(e.g. ~/.warpdata/cache/datasets/.../materialized.parquet)"
                    )
            else:
                typer.echo("Detected multiple Parquet files for this warp dataset:")
                for idx, p in enumerate(parquet_paths, start=1):
                    typer.echo(f"  [{idx}] {p}")
                choice = typer.prompt(
                    "Select index to use (or 0 to enter manually)", default="1"
                )
                try:
                    idx = int(choice)
                except ValueError:
                    idx = 0
                if 1 <= idx <= len(parquet_paths):
                    local_path = str(parquet_paths[idx - 1])
                else:
                    local_path = typer.prompt(
                        "Enter local Parquet path for the main table "
                        "(e.g. ~/.warpdata/cache/datasets/.../materialized.parquet)"
                    )
        else:
            local_path = typer.prompt(
                "Enter local Parquet path for the main table "
                "(e.g. ~/.warpdata/cache/datasets/.../materialized.parquet)"
            )
        # Try to compute approximate size.
        size_info = ""
        try:
            path_obj = Path(local_path).expanduser()
            if path_obj.exists():
                gb = path_obj.stat().st_size / (1024**3)
                size_info = f"# Approx Parquet size: {gb:.2f} GB\n"
        except Exception:
            pass
        hf_id = plan.hf_ids[0]
        typer.echo(
            "\n# Publish plain table to HF and generate schema\n"
            f"{size_info}"
            f"hfschema publish-hf {hf_id} \\\n"
            "  --format parquet \\\n"
            f"  --path {local_path} \\\n"
            "  --split train --private \\\n"
            "  --max-shard-size 3GB\n"
        )

    elif plan.pattern == PlanKind.IMAGE_WITH_IMAGES:
        # Prefer resolving from a dedicated image warp dataset if provided.
        warp_source = getattr(plan, "source_warp_name", None) or plan.warp_name
        parquet_paths = _resolve_warp_parquet_paths(warp_source)
        if parquet_paths:
            if len(parquet_paths) == 1:
                auto = parquet_paths[0]
                typer.echo(f"Detected Parquet file: {auto}")
                if typer.confirm("Use this path?", default=True):
                    local_path = str(auto)
                else:
                    local_path = typer.prompt(
                        "Enter local Parquet path with image_path column "
                        "(e.g. ~/.warpdata/cache/datasets/vision/.../materialized.parquet)"
                    )
            else:
                typer.echo("Detected multiple Parquet files for this warp dataset:")
                for idx, p in enumerate(parquet_paths, start=1):
                    typer.echo(f"  [{idx}] {p}")
                choice = typer.prompt(
                    "Select index to use (or 0 to enter manually)", default="1"
                )
                try:
                    idx = int(choice)
                except ValueError:
                    idx = 0
                if 1 <= idx <= len(parquet_paths):
                    local_path = str(parquet_paths[idx - 1])
                else:
                    local_path = typer.prompt(
                        "Enter local Parquet path with image_path column "
                        "(e.g. ~/.warpdata/cache/datasets/vision/.../materialized.parquet)"
                    )
        else:
            local_path = typer.prompt(
                "Enter local Parquet path with image_path column "
                "(e.g. ~/.warpdata/cache/datasets/vision/.../materialized.parquet)"
            )
        hf_id = plan.hf_ids[0]
        size_info = ""
        try:
            path_obj = Path(local_path).expanduser()
            if path_obj.exists():
                gb = path_obj.stat().st_size / (1024**3)
                size_info = f"# Approx Parquet size (table only): {gb:.2f} GB\n"
        except Exception:
            pass
        typer.echo(
            "\n# Publish image dataset with actual images uploaded to HF\n"
            f"{size_info}"
            f"hfschema publish-hf {hf_id} \\\n"
            "  --format parquet \\\n"
            f"  --images \\\n"
            f"  --path {local_path} \\\n"
            "  --split train --private \\\n"
            "  --max-shard-size 3GB\n"
        )

    elif plan.pattern == PlanKind.TABLE_AND_SIGNALS:
        meta_path = typer.prompt(
            "Enter local Parquet path for metadata table "
            "(e.g. dmt_brains_conditions.parquet or vctk_utterances.parquet)"
        )
        signals_path = typer.prompt(
            "Enter local Parquet path for signals/audio table "
            "(e.g. signals.parquet or vctk_utterances.parquet)"
        )
        meta_id, signals_id = plan.hf_ids
        # Size information for both tables.
        size_meta = ""
        size_sig = ""
        try:
            mp = Path(meta_path).expanduser()
            if mp.exists():
                gb = mp.stat().st_size / (1024**3)
                size_meta = f"# Approx metadata Parquet size: {gb:.2f} GB\n"
        except Exception:
            pass
        try:
            sp = Path(signals_path).expanduser()
            if sp.exists():
                gb = sp.stat().st_size / (1024**3)
                size_sig = f"# Approx signals Parquet size: {gb:.2f} GB\n"
        except Exception:
            pass
        typer.echo(
            "\n# Publish metadata table\n"
            f"{size_meta}"
            f"hfschema publish-hf {meta_id} \\\n"
            "  --format parquet \\\n"
            f"  --path {meta_path} \\\n"
            "  --split train --private \\\n"
            "  --max-shard-size 3GB\n"
        )
        typer.echo(
            "# Publish signals / utterances table\n"
            f"{size_sig}"
            f"hfschema publish-hf {signals_id} \\\n"
            "  --format parquet \\\n"
            f"  --path {signals_path} \\\n"
            "  --split train --private \\\n"
            "  --max-shard-size 3GB\n"
        )

    elif plan.pattern == PlanKind.TABLE_AND_RAW:
        typer.echo(
            "\nThis plan expects a metadata Parquet plus a separate raw-file "
            "HF repo (e.g. BVH files for mocap). For CMU mocap, you can use "
            "the helper script:\n"
            "  bash scripts/publish_cmu_mocap.sh\n"
        )

    elif plan.pattern in (PlanKind.METADATA_ONLY, PlanKind.URL_METADATA_ONLY):
        # Try to resolve Parquet from warp for metadata-only / URL-based tables.
        warp_source = plan.warp_name
        parquet_paths = _resolve_warp_parquet_paths(warp_source)
        prompt_hint = ""
        ws_lower = warp_source.lower()
        if "laion" in ws_lower:
            prompt_hint = "LAION subset parquet"
        elif "sam-3d-body" in ws_lower or "sam_3d_body" in ws_lower:
            prompt_hint = "SAM 3D body metadata parquet"
        else:
            prompt_hint = "metadata parquet"

        if parquet_paths:
            if len(parquet_paths) == 1:
                auto = parquet_paths[0]
                typer.echo(f"Detected Parquet file: {auto}")
                if typer.confirm("Use this path?", default=True):
                    local_path = str(auto)
                else:
                    local_path = typer.prompt(
                        f"Enter local Parquet path for the metadata table "
                        f"(e.g. {prompt_hint})"
                    )
            else:
                typer.echo("Detected multiple Parquet files for this warp dataset:")
                for idx, p in enumerate(parquet_paths, start=1):
                    typer.echo(f"  [{idx}] {p}")
                choice = typer.prompt(
                    "Select index to use (or 0 to enter manually)", default="1"
                )
                try:
                    idx = int(choice)
                except ValueError:
                    idx = 0
                if 1 <= idx <= len(parquet_paths):
                    local_path = str(parquet_paths[idx - 1])
                else:
                    local_path = typer.prompt(
                        f"Enter local Parquet path for the metadata table "
                        f"(e.g. {prompt_hint})"
                    )
        else:
            local_path = typer.prompt(
                f"Enter local Parquet path for the metadata table "
                f"(e.g. {prompt_hint})"
            )
        hf_id = plan.hf_ids[0]
        size_info = ""
        try:
            path_obj = Path(local_path).expanduser()
            if path_obj.exists():
                gb = path_obj.stat().st_size / (1024**3)
                size_info = f"# Approx Parquet size: {gb:.2f} GB\n"
        except Exception:
            pass
        typer.echo(
            "\n# Publish metadata-only table to HF\n"
            f"{size_info}"
            f"hfschema publish-hf {hf_id} \\\n"
            "  --format parquet \\\n"
            "  --no-images \\\n"
            f"  --path {local_path} \\\n"
            "  --split train --private \\\n"
            "  --max-shard-size 3GB\n"
        )

    else:
        typer.echo(
            "\n(No concrete command template for this pattern yet. "
            "Please use `hfschema publish-hf` manually.)"
        )


def _resolve_warp_parquet_paths(warp_name: str) -> list[Path]:
    """
    Best-effort resolution of Parquet paths for a warp dataset.

    Uses warpdata's registry + cache without importing it at module import
    time. Returns an empty list if warpdata is unavailable or the dataset
    cannot be resolved.
    """
    try:
        from warpdata.core.registry import get_registry  # type: ignore[import]
        from warpdata.core.uris import parse_uri  # type: ignore[import]
        from warpdata.core.cache import get_cache  # type: ignore[import]
    except Exception:
        return []

    ds = warp_name.strip()
    if not ds:
        return []

    # Construct a warpdata:// URI if one isn't provided.
    if ds.startswith("warpdata://"):
        uri_str = ds
    elif "/" in ds:
        uri_str = f"warpdata://{ds}"
    else:
        uri_str = f"warpdata://{ds}"

    try:
        uri = parse_uri(uri_str)
        reg = get_registry()
        ver = reg.get_dataset_version(uri.workspace, uri.name, "latest")
        man = reg.get_manifest(uri.workspace, uri.name, ver["version_hash"])
        cache = get_cache()
    except Exception:
        return []

    paths: list[Path] = []
    for res in man.get("resources", []):
        uri_str = res.get("uri")
        if not uri_str:
            continue
        try:
            local_path = cache.get(uri_str)
        except Exception:
            continue
        if str(local_path).endswith(".parquet"):
            paths.append(Path(local_path))

    # Deduplicate while preserving order.
    seen: set[str] = set()
    unique_paths: list[Path] = []
    for p in paths:
        s = str(p)
        if s in seen:
            continue
        seen.add(s)
        unique_paths.append(p)
    return unique_paths


@embeddings_app.command("add")
def embeddings_add(
    space: str = typer.Option(
        ...,
        "--space",
        "-s",
        help="Logical embedding space name (e.g. 'text_gemma').",
    ),
    model: str = typer.Option(
        ...,
        "--model",
        "-m",
        help="Model identifier backing this space.",
    ),
    provider: str | None = typer.Option(
        None,
        "--provider",
        help="Optional provider label (e.g. 'hf', 'openai').",
    ),
    dim: int | None = typer.Option(
        None,
        "--dim",
        help="Optional embedding dimension.",
    ),
    store: Path | None = typer.Option(
        None,
        "--store",
        help="Path to the embeddings configuration file.",
    ),
) -> None:
    """
    Define or update an embedding space that can be linked to schemas.
    """
    store_path = store or default_embeddings_store()
    _add_embedding_space(space, model, provider=provider, dim=dim, store_path=store_path)
    typer.echo(f"Registered embedding space '{space}' with model '{model}'.")


@embeddings_app.command("list")
def embeddings_list(
    store: Path | None = typer.Option(
        None,
        "--store",
        help="Path to the embeddings configuration file.",
    ),
) -> None:
    """
    List configured embedding spaces.
    """
    store_path = store or default_embeddings_store()
    config = load_embeddings_config(store_path)
    if not config.spaces:
        typer.echo("No embedding spaces defined.")
        return
    for name in sorted(config.spaces):
        spec = config.spaces[name]
        model = spec.get("model", "")
        provider = spec.get("provider")
        dim = spec.get("dim")
        parts = [name, f"(model={model}"]
        if provider:
            parts.append(f"provider={provider}")
        if dim is not None:
            parts.append(f"dim={dim}")
        parts.append(")")
        typer.echo(" ".join(parts))


@embeddings_app.command("link-field")
def embeddings_link_field(
    dataset: str = typer.Argument(
        ...,
        help="Dataset id, e.g. 'ag_news' or 'owner/name'.",
    ),
    field: str = typer.Option(
        ...,
        "--field",
        "-f",
        help="Field name within the schema to link.",
    ),
    space: str = typer.Option(
        ...,
        "--space",
        "-s",
        help="Embedding space name to associate with this field.",
    ),
    config: str | None = typer.Option(
        None,
        "--config",
        help="Optional config name for the dataset.",
    ),
    store: Path | None = typer.Option(
        None,
        "--store",
        help="Path to the embeddings configuration file.",
    ),
) -> None:
    """
    Link a schema field to a named embedding space.
    """
    from .naming import schema_key as _schema_key

    store_path = store or default_embeddings_store()
    dataset_key = _schema_key(dataset, config)
    _link_embedding_field(dataset_key, field, space, store_path=store_path)
    typer.echo(f"Linked {dataset_key}.{field} -> {space}")


def _resolve_destination(destination: Path, module_name: str) -> Path:
    destination = Path(destination)
    if destination.exists() and destination.is_dir():
        return destination / module_name
    if destination.suffix:
        return destination
    if destination.exists():
        return destination
    return destination / module_name


@dataclass
class DatasetSelection:
    dataset_id: str
    config: str | None
    split: str | None


@dataclass
class LoaderEntry:
    dataset_id: str
    config: str | None
    split: str
    import_path: str
    schema_class: str


def _parse_dataset_spec(value: str) -> DatasetSelection:
    base, _, split = value.partition("@")
    dataset_id, config = _split_dataset_config(base)
    return DatasetSelection(dataset_id=dataset_id, config=config, split=split or None)


def _split_dataset_config(value: str) -> tuple[str, str | None]:
    if ":" in value:
        ds, cfg = value.split(":", 1)
        return ds.strip(), (cfg.strip() or None)
    return value.strip(), None


def _resolve_schema_selection(dataset: str, config: str | None) -> tuple[str, str | None]:
    registry.discover()
    names = registry.list_names()
    if not names:
        raise typer.Exit(code=1)
    canonical = _match_schema_name(dataset, names)
    dataset_id, existing_config = _split_schema_key(canonical)
    return dataset_id, (config or existing_config)


def _split_schema_key(value: str) -> tuple[str, str | None]:
    if ":" in value:
        dataset_id, config = value.split(":", 1)
        return dataset_id, (config or None)
    return value, None


def _normalize(value: str) -> str:
    return re.sub(r"[^\w]+", "", value.lower())


def _tokenize(value: str) -> list[str]:
    return [token for token in re.split(r"[^\w]+", value.lower()) if token]


def _match_schema_name(query: str, names: list[str]) -> str:
    query = query.strip()
    if not query:
        raise typer.BadParameter("Dataset name cannot be empty.")

    normalized_query = query.lower()
    names_sorted = sorted(names)
    lookup = {name.lower(): name for name in names_sorted}
    if normalized_query in lookup:
        return lookup[normalized_query]

    split_query, split_config = _split_schema_key(query)
    dataset_only = split_query.lower()
    dataset_lookup = {name.split(":", 1)[0].lower(): name for name in names_sorted}
    if normalized_query in dataset_lookup:
        return dataset_lookup[normalized_query]
    if dataset_only in dataset_lookup:
        candidate = dataset_lookup[dataset_only]
        if split_config:
            dataset_id, _ = _split_schema_key(candidate)
            return f"{dataset_id}:{split_config}"
        return candidate

    # substring match
    for name in names_sorted:
        if normalized_query in name.lower():
            return name

    query_tokens = _tokenize(query)
    if query_tokens:
        for name in names_sorted:
            candidate_tokens = _tokenize(name)
            if all(
                any(q in cand for cand in candidate_tokens) for q in query_tokens
            ):
                return name

    normalized = _normalize(query)
    for name in names_sorted:
        if normalized in _normalize(name):
            return name

    # fallback to first entry
    return names_sorted[0]


def _build_loader_entry(dataset_id: str, config: str | None, split: str) -> LoaderEntry:
    module_stem = module_filename(dataset_id, config).rsplit(".", 1)[0]
    import_path = f"schemas_hf.datasets.{module_stem}"
    module = importlib.import_module(import_path)
    schema_cls = _first_schema_class(module)
    if schema_cls is None:
        raise RuntimeError(f"No DatasetSchema subclass found in {import_path}")
    return LoaderEntry(
        dataset_id=dataset_id,
        config=config,
        split=split,
        import_path=import_path,
        schema_class=schema_cls.__name__,
    )


def _first_schema_class(module):
    from schemas_hf.base import DatasetSchema
    import inspect

    for _, obj in inspect.getmembers(module, inspect.isclass):
        if issubclass(obj, DatasetSchema) and obj is not DatasetSchema:
            return obj
    return None


def _slugify_name(name: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9]+", "_", name.strip()).strip("_")
    return slug or "loader"


def _render_loader_template(loader_name: str, entries: list[LoaderEntry], streaming: bool) -> str:
    import textwrap

    import_lines = sorted(
        {f"from {entry.import_path} import {entry.schema_class}" for entry in entries}
    )
    dataset_lines = []
    for entry in entries:
        config_repr = "None" if entry.config is None else repr(entry.config)
        dataset_lines.append(
            f'    {{"id": "{entry.dataset_id}", "config": {config_repr}, '
            f'"split": "{entry.split}", "schema": {entry.schema_class}}},'
        )

    lines = [
        "# Auto-generated by hfschema. Edit as needed.",
        "from typing import Dict, Iterator",
        "",
        "from datasets import load_dataset",
        "",
    ]
    lines.extend(import_lines)
    if import_lines:
        lines.append("")
    lines.append("DATASETS = [")
    lines.extend(dataset_lines)
    lines.append("]")
    lines.append("")
    lines.append("")
    lines.append(f"def iter_examples(streaming: bool = {streaming}) -> Iterator[Dict]:")
    lines.append('    """Yield validated records from each configured dataset."""')
    lines.append("    for spec in DATASETS:")
    lines.append("        ds = load_dataset(")
    lines.append('            spec["id"],')
    lines.append('            spec["config"],')
    lines.append('            split=spec["split"],')
    lines.append("            streaming=streaming,")
    lines.append("        )")
    lines.append("        for row in ds:")
    lines.append('            record = spec["schema"].from_example(row)')
    lines.append("            yield {")
    lines.append('                "dataset": spec["id"],')
    lines.append('                "config": spec["config"],')
    lines.append('                "split": spec["split"],')
    lines.append('                "record": record,')
    lines.append("            }")
    lines.append("")
    lines.append("")
    lines.append(
        f"def batched_examples(batch_size: int = 32, streaming: bool = {streaming}) -> Iterator[list]:"
    )
    lines.append('    """Yield batches of validated examples."""')
    lines.append("    batch: list = []")
    lines.append("    for example in iter_examples(streaming=streaming):")
    lines.append("        batch.append(example)")
    lines.append("        if len(batch) == batch_size:")
    lines.append("            yield batch")
    lines.append("            batch = []")
    lines.append("    if batch:")
    lines.append("        yield batch")
    lines.append("")

    return textwrap.dedent("\n".join(lines))


def _echo_dataset_error(dataset_id: str, config: str | None, exc: Exception) -> None:
    target = f"{dataset_id}:{config}" if config else dataset_id
    typer.echo(f"Unable to load metadata for '{target}'.", err=True)
    typer.echo(f"Reason: {exc}", err=True)
    typer.echo(
        "Check that the dataset/config exists, you have access (hf-cli login), "
        "and your network/cache is available.",
        err=True,
    )


if __name__ == "__main__":
    app()
