"""
hegm._hook -- Generic PEP-451 import hook.

Installs a one-shot ``MetaPathFinder`` for each registered backend so that
the framework module (e.g. ``torch``, ``tensorflow``) is intercepted on
first import and the backend's ``apply_patches`` callback is invoked.
"""

import importlib
import importlib.abc
import importlib.machinery
import importlib.util
import sys
from typing import Any, Callable, Optional

from hegm._config import log

# ---------------------------------------------------------------------------
# Loader / Finder (parameterised by *target_module* and *callback*)
# ---------------------------------------------------------------------------


class _HeGMLoader(importlib.abc.Loader):
    """Wraps a real loader to invoke *callback* after ``exec_module``."""

    def __init__(
        self,
        original_loader: importlib.abc.Loader,
        callback: Callable[[Any], None],
    ) -> None:
        self._original = original_loader
        self._callback = callback

    def create_module(
        self, spec: importlib.machinery.ModuleSpec
    ) -> Optional[Any]:
        if hasattr(self._original, "create_module"):
            return self._original.create_module(spec)
        return None

    def exec_module(self, module: Any) -> None:
        self._original.exec_module(module)
        self._callback(module)


class _HeGMFinder(importlib.abc.MetaPathFinder):
    """One-shot finder that intercepts ``import <target_module>``.

    1. Remove ourselves from ``sys.meta_path``.
    2. Locate the *real* ``ModuleSpec``.
    3. Swap its loader with ``_HeGMLoader``.
    4. Return the modified spec.
    """

    def __init__(
        self,
        target_module: str,
        callback: Callable[[Any], None],
    ) -> None:
        self._target = target_module
        self._callback = callback

    def find_spec(
        self,
        fullname: str,
        path: Any = None,
        target: Any = None,
    ) -> Optional[importlib.machinery.ModuleSpec]:
        if fullname != self._target:
            return None

        try:
            sys.meta_path.remove(self)
        except ValueError:
            return None  # already removed (concurrent import race)

        real_spec = importlib.util.find_spec(self._target)
        if real_spec is None:
            return None  # framework not installed

        real_spec.loader = _HeGMLoader(real_spec.loader, self._callback)  # type: ignore[assignment]
        return real_spec


# ---------------------------------------------------------------------------
# Public helper
# ---------------------------------------------------------------------------


def install_hook(target_module: str, callback: Callable[[Any], None]) -> None:
    """Install a one-shot import hook for *target_module*.

    If *target_module* is already imported, *callback* is called immediately.
    """
    if target_module in sys.modules:
        log(f"{target_module} already imported -- applying patches directly.",
            level="DEBUG")
        callback(sys.modules[target_module])
    else:
        sys.meta_path.insert(0, _HeGMFinder(target_module, callback))
