import sys
from importlib import import_module

__all__ = ["LazyImport", "LazyImportErrorProxy" "is_imported", "import_common"]


def _try_import(pkg):
    try:
        _module = __import__(str(pkg))
        return _module
    except ImportError:
        return None


def is_imported(module_name: str):
    ans = module_name in sys.modules
    print(f"{module_name}{''if ans else ' not'} in sys.modules")
    return ans


class LazyImport:
    """LazyImport

    Examples:
        # Case 1: Standard module import
        np = LazyImport("numpy")
        x = np.array([1, 2, 3])

        # Case 2: Direct class import
        Path = LazyImport("pathlib", object_name="Path")
        p = Path("file.txt")
    """

    def __init__(self, module_name, object_name=None):
        self.module_name = module_name
        self.object_name = object_name
        self._target = None

    def _load_target(self):
        module = import_module(self.module_name)
        if not self.object_name:
            self._target = module
        else:
            self._target = getattr(module, self.object_name)

    def __getattr__(self, name):
        if self._target is None:
            self._load_target()
        return getattr(self._target, name)

    def __getitem__(self, key):
        if self._target is None:
            self._load_target()
        return self._target.__getitem__(key)

    def __call__(self, *args, **kwargs):
        if self._target is None:
            self._load_target()
        return self._target(*args, **kwargs)


class LazyImportErrorProxy:
    """
    A proxy class that delays raising an ImportError until the object is actually
    called or instantiated. This allows the code to run as long as the missing
    library's features are not invoked.
    """

    def __init__(self, module_name: str, msg: str):
        self._module_name = module_name
        self._msg = msg

    def __getattr__(self, name):
        # Support chained attribute access (e.g., rich.console.Console)
        # By returning 'self', any sub-attribute will also be a proxy instance.
        return self

    def __call__(self, *args, **kwargs):
        # Raise the error only when the code tries to actually USE the object
        raise ImportError(
            f"Error: The '{self._module_name}' library is required but not installed. "
            f"Please install it by running 'pip install {self._module_name}'."
            f" {self._msg}"
        )

    def __repr__(self):
        return f"<Missing module proxy for '{self._module_name}'>"


def import_common(g=None):
    if g is None:
        g = globals()
    g["torch"] = LazyImport("torch")
    g["F"] = LazyImport("torch.nn.functional")
    g["np"] = LazyImport("numpy")
    g["pd"] = LazyImport("pandas")
    g["plt"] = LazyImport("matplotlib", "pyplot")

    g["Path"] = LazyImport("pathlib", "Path")
    g["os"] = LazyImport("os")
    g["sys"] = LazyImport("sys")
    g["time"] = LazyImport("time")
    g["json"] = LazyImport("json")
    g["pickle"] = LazyImport("pickle")
    g["tqdm"] = LazyImport("tqdm", "tqdm")
    g["trange"] = LazyImport("tqdm", "trange")

    # typing
    g["Any"] = LazyImport("typing", "Any")
    g["Dict"] = LazyImport("typing", "Dict")
    g["Iterable"] = LazyImport("typing", "Iterable")
    g["List"] = LazyImport("typing", "List")
    g["Optional"] = LazyImport("typing", "Optional")
    g["Sequence"] = LazyImport("typing", "Sequence")
    g["Tuple"] = LazyImport("typing", "Tuple")
    g["Union"] = LazyImport("typing", "Union")
