# Use CSafeFile if available
from typing import TYPE_CHECKING
from typing import Any
from typing import Dict
from typing import Iterable
from typing import Tuple

if TYPE_CHECKING:
    from yaml import SafeLoader
else:
    try:
        from yaml import CSafeLoader as SafeLoader
    except ImportError:
        from yaml import SafeLoader


__all__ = [
    "SafeLoader",
]


class LimitedSafeLoader(type):
    """Meta YAML loader that skips the resolution of the specified YAML tags."""

    def __new__(
        cls,
        name: str,
        bases: Tuple[type, ...],
        namespace: Dict[str, Any],
        exclude_resolvers: Iterable[str],
    ) -> "LimitedSafeLoader":
        exclude_resolvers = set(exclude_resolvers)
        implicit_resolvers = {
            key: [
                (tag, regex)
                for tag, regex in mappings
                if tag not in exclude_resolvers
            ]
            for key, mappings in SafeLoader.yaml_implicit_resolvers.items()
        }
        return super().__new__(
            cls,
            name,
            (SafeLoader, *bases),
            {**namespace, "yaml_implicit_resolvers": implicit_resolvers},
        )


class JsonschemaSafeLoader(
    metaclass=LimitedSafeLoader,
    exclude_resolvers={"tag:yaml.org,2002:timestamp"},
):
    """A safe YAML loader that leaves timestamps as strings."""