from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Iterable, Iterator, Set, Tuple

from .cache import (
    DatasetConfigRef,
    DatasetInfoFile,
    iter_dataset_cache_configs,
    iter_dataset_info_files,
    iter_repo_ids,
)
from .codegen import SchemaRenderer
from .naming import module_filename

def _standardize_dataset_id(value: str) -> str:
    return "".join(ch.lower() for ch in value if ch.isalnum())

AUTO_HEADER = "# Auto-generated by schemas-hf. Do not edit manually.\n\n"
DEFAULT_DATASET_CACHE = Path.home() / ".cache" / "huggingface" / "datasets"
LOCAL_DATASETS = {"json", "text"}


def generate_from_cache(
    cache_dir: Path,
    output_dir: Path,
    *,
    dataset_cache_dir: Path | None = None,
    dry_run: bool = False,
    only: Iterable[str] | None = None,
    skip_existing: bool = False,
) -> list[Path]:
    """Generate schema modules from a Hugging Face cache directory."""
    cache_dir = Path(cache_dir)
    output_dir = Path(output_dir)
    dataset_cache_dir = Path(dataset_cache_dir or DEFAULT_DATASET_CACHE)

    dataset_infos, hub_dataset_ids = _collect_dataset_infos(cache_dir)
    hub_lookup = {_standardize_dataset_id(ds): ds for ds in hub_dataset_ids}
    targets = set(dataset_infos.keys())

    dataset_filter, config_filter = _parse_only_filter(only)

    for ref in iter_dataset_cache_configs(dataset_cache_dir):
        normalized = _normalize_config_name(ref.config_name)
        dataset_id = hub_lookup.get(_standardize_dataset_id(ref.dataset_id), ref.dataset_id)
        if dataset_id in LOCAL_DATASETS and dataset_id not in dataset_filter:
            if (dataset_id, normalized) not in config_filter:
                continue
        if hub_dataset_ids and dataset_id not in hub_dataset_ids:
            if (
                dataset_id not in dataset_filter
                and (dataset_id, normalized) not in config_filter
            ):
                continue
        targets.add((dataset_id, normalized))

    targets = _filter_targets(targets, dataset_filter, config_filter)

    generated: list[Path] = []
    seen: set[tuple[str, str | None]] = set()

    for dataset_id, config_name in sorted(
        targets, key=lambda item: (item[0], item[1] or "")
    ):
        key = (dataset_id, config_name)
        if key in seen:
            continue
        seen.add(key)

        module_path = output_dir / module_filename(dataset_id, config_name)

        # When skip_existing is enabled, we avoid calling the HF builder if a
        # schema module for this dataset/config already exists. This keeps
        # subsequent `sync` runs fast once schemas are generated.
        if skip_existing and module_path.exists() and not dry_run:
            generated.append(module_path)
            continue

        info = dataset_infos.get(key)
        if info is None:
            info = _load_from_builder(dataset_id, config_name)
            dataset_infos[key] = info

        renderer = SchemaRenderer(dataset_id, config_name)
        contents = AUTO_HEADER + renderer.render(info)

        if not dry_run:
            module_path.parent.mkdir(parents=True, exist_ok=True)
            module_path.write_text(contents)
            generated.append(module_path)

    return generated


def _collect_dataset_infos(
    cache_dir: Path,
) -> tuple[Dict[Tuple[str, str | None], dict], Set[str]]:
    infos: Dict[Tuple[str, str | None], dict] = {}
    dataset_ids: Set[str] = set(iter_repo_ids(cache_dir))
    for info_file in iter_dataset_info_files(cache_dir):
        dataset_ids.add(info_file.repo_id)
        for config_name, info in _parse_dataset_info(info_file):
            infos[(info_file.repo_id, config_name)] = info
    return infos, dataset_ids


def _parse_dataset_info(info_file: DatasetInfoFile) -> Iterator[tuple[str | None, dict]]:
    try:
        text = info_file.path.read_text().strip()
    except FileNotFoundError:
        return iter(())
    if not text:
        return iter(())
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        return iter(())

    def _iter_single(data: dict, config_name: str | None):
        yield (_normalize_config_name(config_name), data)

    if "features" in data:
        return _iter_single(data, data.get("config_name"))

    output: list[tuple[str | None, dict]] = []
    for config_name, value in data.items():
        if isinstance(value, dict) and "features" in value:
            output.append((_normalize_config_name(config_name), value))
    return iter(output)


def _parse_only_filter(
    only: Iterable[str] | None,
) -> tuple[Set[str], Set[tuple[str, str | None]]]:
    if not only:
        return set(), set()
    dataset_filter: Set[str] = set()
    config_filter: Set[tuple[str, str | None]] = set()
    for item in only:
        if ":" in item:
            dataset_id, config_name = item.split(":", 1)
            config_filter.add((dataset_id, _normalize_config_name(config_name)))
        else:
            dataset_filter.add(item)
    return dataset_filter, config_filter


def _filter_targets(
    targets: Set[tuple[str, str | None]],
    dataset_filter: Set[str],
    config_filter: Set[tuple[str, str | None]],
) -> Set[tuple[str, str | None]]:
    if not dataset_filter and not config_filter:
        return targets

    selected: Set[tuple[str, str | None]] = set()
    for dataset_id, config_name in targets:
        if dataset_id in dataset_filter or (dataset_id, config_name) in config_filter:
            selected.add((dataset_id, config_name))

    for dataset_id in dataset_filter:
        if not any(item[0] == dataset_id for item in selected):
            selected.add((dataset_id, None))

    selected |= config_filter
    return selected


def _normalize_config_name(name: str | None) -> str | None:
    if not name or name == "default":
        return None
    return name


def _load_from_builder(dataset_id: str, config_name: str | None) -> dict:
    from datasets import load_dataset_builder

    base_kwargs = {"name": config_name} if config_name else {}
    try:
        # Prefer enabling remote code for script-based datasets, but fall back
        # gracefully for packaged datasets or older `datasets` versions where
        # the BuilderConfig does not accept `trust_remote_code`.
        try:
            builder = load_dataset_builder(
                dataset_id,
                trust_remote_code=True,
                **base_kwargs,
            )
        except ValueError as inner_exc:
            message = str(inner_exc)
            if "trust_remote_code" in message and "doesn't have a" in message:
                # Retry without the extra kwarg; this covers packaged datasets
                # like CSV as well as older `datasets` versions.
                builder = load_dataset_builder(dataset_id, **base_kwargs)
            else:
                raise
    except Exception as exc:  # pragma: no cover - depends on remote datasets
        raise RuntimeError(
            f"Unable to load dataset info for {dataset_id}"
            + (f":{config_name}" if config_name else "")
        ) from exc
    features = builder.info.features.to_dict() if builder.info.features else {}
    return {"features": features}


__all__ = ["generate_from_cache"]
