from __future__ import annotations

from dataclasses import is_dataclass
from importlib import import_module
from pkgutil import iter_modules
from typing import Iterable

from .base import DatasetSchema
from .naming import schema_key


class SchemaRegistry:
    """Keeps track of available dataset schemas."""

    def __init__(self) -> None:
        self._schemas: dict[tuple[str, str | None], type[DatasetSchema]] = {}

    def register(self, schema_cls: type[DatasetSchema]) -> None:
        if not is_dataclass(schema_cls):
            raise TypeError("Only dataclass-based schemas can be registered.")
        if not issubclass(schema_cls, DatasetSchema):
            raise TypeError("Schema must inherit from DatasetSchema.")

        key = (schema_cls.dataset_name(), schema_cls.config_name())
        if key in self._schemas:
            raise ValueError(f"Schema '{schema_cls.dataset_name()}' already registered.")
        self._schemas[key] = schema_cls

    def get(self, name: str) -> type[DatasetSchema]:
        for (dataset_name, config_name), schema in self._schemas.items():
            key = schema_key(dataset_name, config_name)
            if key == name:
                return schema
        raise KeyError(name)

    def list_names(self) -> list[str]:
        return sorted(
            schema_key(dataset_name, config_name)
            for dataset_name, config_name in self._schemas
        )

    def discover(self, package: str = "schemas_hf.datasets") -> None:
        """Import all modules under the package so decorators can register schemas."""
        module = import_module(package)
        if getattr(module, "__path__", None) is None:
            return

        for finder, name, ispkg in iter_modules(module.__path__, prefix=f"{package}."):
            try:
                import_module(name)
            except Exception:
                # If a generated schema module is syntactically invalid or
                # otherwise fails to import, skip it. This keeps the registry
                # usable even when individual auto-generated files are broken.
                continue


registry = SchemaRegistry()


def register_schema(cls: type[DatasetSchema]) -> type[DatasetSchema]:
    """Class decorator to register schemas on import."""
    registry.register(cls)
    return cls
