# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Module registry for Fuel compiler.

The module registry implements the APIs that are called by the
translator in response to import statements.  The ModuleRegistry class
defines methods importmodule() and importfrom(), which are called for
a plain import statement and a from-import statement, respectively.
The LoadError exception is raised by either of those modules if the
requested module cannot be found.  The Module class represents a
module object as returned by the import methods; it defines a lookup()
method that returns a symbol definition.

The ModuleRegistry class here loads module contents from 'compiled'
Fuel files, which have a .fuelc filename extension.  There's a
subclass defined in the compile module which overrides much of its
implementation to cause on-demand compilation.

"""

# Python imports
import os
import pickle

# Fuel compiler imports
from sigma.fuelc import symtab
from sigma.fuelc import typesys

# Sigma imports
from sigma.util.fuelset import FuelSet


class _Py2Set(FuelSet):
    """Stand-in for Python 2's ``sets.Set`` during unpickling.

    Python 2's ``sets.Set`` stored its elements as keys of an internal
    ``_data`` dict.  The pickle stream uses ``__setstate__`` to pass
    this dict.  Python 3's built-in ``set`` does not accept a dict in
    ``__setstate__``, so this subclass intercepts it and populates the
    set from the dict keys.
    """

    def __setstate__(self, state):
        if isinstance(state, dict):
            self.update(state.keys())
        elif state is not None:
            self.update(state)


class _Py2ImmutableSet(frozenset):
    """Stand-in for Python 2's ``sets.ImmutableSet`` during unpickling."""

    def __new__(cls, iterable=()):
        return super().__new__(cls, iterable)

    def __setstate__(self, state):
        # frozenset is immutable -- we can't add elements after creation.
        # Python 2 ImmutableSet pickled as (constructor(), __setstate__(data))
        # but frozenset must receive its elements at __new__ time.
        # If we get here, the elements were not passed to __new__,
        # so we need to work around this by rebuilding.
        pass


class _Py2ImmutableSetBuilder:
    """Factory for unpickling Python 2's ``sets.ImmutableSet``.

    The pickle protocol creates an empty instance via __new__/__init__ then
    calls __setstate__ with the data dict.  Since frozenset is immutable,
    we collect the state and return a proper frozenset via __reduce__.
    """

    def __init__(self):
        self._elements = ()

    def __setstate__(self, state):
        if isinstance(state, dict):
            self._elements = tuple(state.keys())
        elif state is not None:
            self._elements = tuple(state)

    def __reduce__(self):
        return (frozenset, (self._elements,))


class Py2Unpickler(pickle.Unpickler):
    """Unpickler that maps Python 2 types to their Python 3 equivalents.

    The .fuelc pickle files were originally created under Python 2 and may
    contain references to the now-removed ``sets`` module (Set, ImmutableSet).
    This subclass transparently replaces them with ``FuelSet`` /
    ``frozenset``.
    """

    def find_class(self, module, name):
        if module == "sets":
            if name == "Set":
                return _Py2Set
            if name in ("ImmutableSet", "BaseSet"):
                return _Py2ImmutableSetBuilder
        if module.startswith("elemental."):
            module = "sigma." + module[len("elemental."):]
        return super().find_class(module, name)


class DuplicateModule(Exception):
    "Exception raised for duplicate module definitions."

class LoadError(Exception):
    "Exception raised for problems loading a module."

class Module(object):

    """A wrapper for imported module's symbol table.

    The symbol table must be passed in to the constructor; there is no
    API to define new symbols.

    Informally, these attributes may be used:

    name       the module name; may contain dots
    symbols    a dict mapping names to SymbolDefinition instances
               (see the symtab module)
    filename   the compiled Fuel filename (with .fuelc extension)

    """

    def __init__(self, name, symbols, filename):
        """Constructor.

        Arguments:

        name       the module name; may contain dots
        symbols    a dict mapping names to SymbolDefinition instances
        filename   the compiled Fuel filename (with .fuelc extension)

        """
        assert isinstance(symbols, dict), symbols
        self.name = name
        self.symbols = symbols
        self.filename = filename

    def lookup(self, name):
        """Look up a symbol name in the module.

        Given a symbol name, return a SymbolDefinition instance that
        defines it, or None if the name is not defined.

        """
        return self.symbols.get(name)

def defaultfuelpath():
    """Return the default Fuel search path, a list of directory paths.

    If $FUELPATH is given in the environment, use that (using the
    system standard search path separator, which is ':' on Unix
    and ';' on Windows).  If $FUELPATH is not given or empty, the
    default search path consists of a single directory, the
    directory represented by //src/server/src/fuel/ in the
    Perforce tree.

    """
    fp = os.getenv("FUELPATH")
    if fp:
        return fp.split(os.pathsep)
    here = os.path.realpath(os.path.dirname(__file__))
    ppparent = os.path.dirname(os.path.dirname(os.path.dirname(here)))
    fueldir = os.path.join(ppparent, "lib", "fuel")
    assert os.path.isdir(fueldir), fueldir
    return [fueldir]

class ModuleRegistry(object):

    """The API used by the translator for handling imports.

    This provides a default implementation that looks for previously
    compiled .fuelc files.

    """

    def __init__(self, fuelpath=None):
        """Constructor.

        The optional argument is the Fuel search path: a list of
        directory pathnames where to look for .fuelc files.  If this
        is not given (or explicitly None), a default Fuel search path
        is constructed by the defaultfuelpath() method.

        """
        if fuelpath is None:
            fuelpath = self.defaultfuelpath()
        self.setfuelpath(fuelpath)
        self.modules = {}

    def setfuelpath(self, fuelpath):
        """Set the fuel path."""
        self.fuelpath = fuelpath
        self.realfuelpath = {}
        for dir in self.fuelpath:
            self.realfuelpath[dir] = os.path.normcase(os.path.realpath(dir))

    defaultfuelpath = staticmethod(defaultfuelpath)

    def getpackagename(self, filename):
        """Map a File filename to a Fuel package name.

        The Fuel filename must end in .fuel.

        The return value is a string giving the package name as a
        dotted name.  For example, if filename == 'X/foo/bar/blah.fuel'
        where X is one of the directories on the Fuel search path, the
        return value would be 'foo.bar'.

        If the file is found in one of the toplevel Fuel path
        directories, or not at all, None is returned.

        """
        head, tail = os.path.split(filename)
        base, ext = os.path.splitext(tail)
        assert ext.lower() == ".fuel", ext
        realhead = os.path.normcase(os.path.realpath(head))
        for dir, realdir in self.realfuelpath.items():
            rh = realhead
            mod = []
            while rh:
                if rh == realdir:
                    if not mod:
                        break
                    return ".".join(mod)
                rh, rest = os.path.split(rh)
                if not rest:
                    break
                mod.insert(0, rest)
        return None # Toplevel module

    def lookup(self, name):
        """Look up a module by name.

        Return a SymbolDefinition object, or None if there is no
        module with the given name.

        Module names may contain dots; e.g. the file x/y/z.fuel
        (somewhere under $FUELPATH) would define a module with the
        name 'x.y.z'.  This module can be used in two ways: 'from x.y
        import z' imports the module itself, while 'from x.y.z import
        f' imports the object f defined in the module.

        """
        return self.modules.get(name)

    def define(self, name, modobj):
        """Define a new module.

        The name must be a module name, possibly containing dots (see
        lookup()).  The module object must be a Module instance.  The
        module must not already be defined; DuplicateModule will be
        raised if it is.

        Return the SymbolDefinition object inserted into self.modules.

        """
        assert isinstance(modobj, Module), modobj
        if name in self.modules:
            raise DuplicateModule(name)
        defn = symtab.SymbolDefinition(name, None, typesys.BasicModule, modobj)
        self.modules[name] = defn
        return defn

    def findmodule(self, name, suffixes=[".fuelc"]):
        """Search the fuel search path for a file for the named module.

        The name argument specifies a module name, possibly containing
        dots.

        The optional suffixes argument is a list of filename
        extensions; it defaults to ['.fuelc'].

        The return value is the full pathname of the file found, or
        None if no match was found.

        Example: if the name argument is 'x.y.z' and the suffixes
        argument is not given, this looks for a file named x/y/z.fuelc
        along the fuel search path.  If the suffixes list is given as
        ['.fuel', '.fuelc'] instead, this looks for files named
        x/y/z.fuel or x/y/z.fuelc along the fuel search path.

        """
        file = os.path.join(*name.split("."))
        files = [file + suffix for suffix in suffixes]
        for dir in self.fuelpath:
            for file in files:
                full = os.path.join(dir, file)
                if os.path.exists(full):
                    return full
        return None

    def loadmodule(self, name, filename):
        """Load the symbol table for the named module from the given file.

        Equivalent to self.loadfrompickle(name, self.loadpickle(filename)).

        (The code is factored in this way so that the
        CompilationManager subclass can extract the dependency data
        from the pickle without having to read the pickle file twice.)

        """
        data = self.loadpickle(filename)
        return self.loadfrompickle(name, data)

    def loadpickle(self, filename):
        """Read pickle data for a compiled module from a file.

        The filename is assumed to be a .fuelc file.  A binary pickle
        is read from it.  This pickle must be a dict and it must have
        a key 'symbols'; this provides a minimal amount of sanity
        checking.  The dict is returned.

        """
        f = open(filename, "rb")
        try:
            data = Py2Unpickler(f).load()
        finally:
            f.close()
        if not isinstance(data, dict):
            raise LoadError("%s: pickle is not a dict" % filename)
        if "symbols" not in data:
            raise LoadError("%s: key 'symbols' not found" % filename)
        data["filename"] = filename # So loadfrompickle() can set this
        return data

    def loadfrompickle(self, name, data):
        """Load a module from pickle data read by loadpickle().

        The name argument is the module name; the data argument is the
        data read from the pickle by loadpickle().  This creates a
        Module object from the symbol table contained in the pickle
        under key 'symbols', and defines that module using
        self.define().

        """
        symbols = data["symbols"]
        modobj = Module(name, symbols, data["filename"])
        defn = self.define(name, modobj)
        return defn

    def importmodule(self, name):
        """Import a module by name.

        This is called by the translator in response to 'import name'.
        In this case the name contains no dots, although the routine
        will also work with dotted names.

        The return value is a SymbolDefinition object.

        If the module could not be loaded for any reason, LoadError is
        raised.  In this case, error messages may already have been
        written to stderr by the parser or the translator if they
        rejected the module; but if the problem was being able to open
        or read the file, no errors have been reported yet and the
        error message is included in the LoadError excetion.

        """
        defn = self.tryimportmodule(name)
        if defn is None:
            raise LoadError("can't find module %s" % name)
        return defn

    def tryimportmodule(self, name):
        """Try to import a module by name.

        This is like importmodule(name), but returns None if the file
        could not be found, rather than raising LoadError in that
        case.  (Other errors during loading still raise LoadError.)

        """
        defn = self.lookup(name)
        if defn is None:
            filename = self.findmodule(name)
            if filename is not None:
                defn = self.loadmodule(name, filename)
                assert defn is not None
        return defn

    def importfrom(self, prefix, name):
        """Import a module in response to a from-import statement.

        The module is specified by a prefix and a name.  The prefix is
        a possibly dotted name; the name is an undotted name.

        The return value is a SymbolDefinition object.

        The prefix should not begin with a dot, even though the
        from-import syntax allows the prefix to begin with a dot; the
        translator should convert such a relative prefix to an
        absolute package name before calling importfrom().

        If the prefix can be seen as both a package and a module, we
        prefer the interpretation as a package.

        In either case, the module may either be loaded from the
        filesystem using loadmodule or we may return a previously
        loaded module (or an object it defines).

        XXX It's possible that the long and short form are both used;
        for example, it's possible that x/y.fuel and x/y/z.fuel both
        exist, and both are loaded.  This can be forced e.g. by the
        following two Fuel statements:

            from x import y
            from x.y import z

        Does this mean that the Fuel syntax should be changed to
        disambiguate this???

        """

        # The prefix should not be relative (the caller should absolutize it):
        assert not prefix.startswith("."), prefix

        # Try to import the long version
        defn = self.tryimportmodule(prefix + "." + name)
        if defn is not None:
            return defn

        # Import the short version and look up name in the result
        defn = self.importmodule(prefix)
        assert defn is not None
        assert isinstance(defn.value, Module)
        defn = defn.value.lookup(name)
        if defn is not None:
            return defn
        raise LoadError("module %s doesn't define %s" % (prefix, name))
