# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Fuel compiler implementation.

Exported class:

  CompilationManager -- manage imports, compilations and dependencies
"""

# Python imports
import os
import sys
import logging
import linecache
import pickle

# Fuel compiler imports
from sigma.fuelc import parse
from sigma.fuelc import driver
from sigma.fuelc import modreg
from sigma.fuelc import symtab
from sigma.fuelc import typesys
from sigma.fuelc import translate
from sigma.util.errors import plural

# Global objects
logger = logging.getLogger("fuelc.compile")

# Standard header added to all compiled output
PROLOGUE = r"""# *** DO NOT EDIT!!!  This is Fuel compiler output. ***
from sigma.util.fuelset import FuelSet
from sigma.util.fuelobjs import FuelStruct
from sigma.runtime.rtutils import *
if __name__ != "fuel.stdlib_fuel":
    from fuel.stdlib_fuel import *

"""

# Standard footer added to non-main compiled output
REG_EPILOGUE = r"""
# This module cannot be used as a main program
if __name__ == '__main__':
    import sys
    sys.stderr.write('ERROR: This module cannot be used as a main program\n')
    sys.exit(2)
"""

# Special footer added to compiled output that has a main() entry point
MAIN_EPILOGUE = r"""
# Main program to start execution at the default entry point
if __name__ == '__main__':
    from sigma.runtime.rtlib import fuel_main
    fuel_main(main)
"""

def getmtime(filename):
    """Helper function to return a file's mtime.

    Return 0 if there was an error.

    Log a debug message.

    """
    try:
        mtime = os.path.getmtime(filename)
    except os.error as err:
        logger.debug("can't get mtime for %s: %s", filename, err)
        return 0
    else:
        logger.debug("mtime for %s -> %s", filename, mtime)
        return mtime

class CompilationManager(modreg.ModuleRegistry):

    """Manage compilations and imports.

    This extends some methods of modreg.ModuleRegistry to ensure that
    when loading a module, the module is first compiled if the .fuelc
    file is out of date.  This also adds new functionality to manage a
    compilations and dependencies.

    The proper call order is:

    - constructor
    - one or more calls to addfile()
    - compileall()
    - one or more getstatus() calls (to check on compilation success)

    For better or for worse, the fuelpath attribute may be used
    read-only to access the Fuel module search path passed in or
    calculated.

    """

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

        Arguments should be given as keyword arguments only:

        fuelpath -- if specified and not None, the search path for
                    Fuel modules.  Otherwise, a default path is chosen
                    by the base class (or by overriding the
                    defaultfuelpath() method).
        force    -- if specified and true, force recompilation of all
                    Fuel code, even if only referenced via import.
        update   -- if specified and true, only compile Fuel modules
                    if the compiler output is found out of date.

        At most one of force and update should be true; by default,
        Fuel files specified explicitly through addfile() are always
        compiled, but dependencies referenced via import are compiled
        only if out of date.

        """
        if force and update:
            raise ValueError("force and update shouldn't both be true")
        modreg.ModuleRegistry.__init__(self, fuelpath)
        self.force = force
        self.update = update
        # A mapping of *normalized* filenames (os.path.realpath()) to
        # status codes: "todo", "busy", "done", "error":
        self.filestatus = {}
        # A mapping of module names to normalized filenames:
        self.savedpaths = {}
        # A mapping of (prefix, name) pairs to either prefix or
        # prefix + "." + name, indicating how "from prefix import name"
        # is satisfied (for thorough dependency checking)
        self.fromimports = {}

    def addfile(self, filename):
        """Add a file for compilation.

        The argument should be a .fuel filename.  If it doesn't exist,
        an IOError will be raised later during compilation.

        The return value is the filename, normalized by
        os.path.realpath(); this can be used to pass to getstatus().

        The file's status is set to 'todo' when it is first added.
        Subsequent calls to add the same file have no effect.

        """
        assert filename.endswith(".fuel"), filename
        filename = os.path.realpath(filename)
        assert filename.endswith(".fuel"), filename
        self.filestatus.setdefault(filename, "todo")
        return filename

    def compileall(self):
        """Compile all files specified by addfile() calls.

        Compilation may happen out of order to satisfy module
        dependencies.  The force and update flags are respected.
        """
        self.scan_stdlib()
        done = {}
        while True:
            for filename, status in list(self.filestatus.items()):
                if status == "todo" and filename not in done:
                    done[filename] = None
                    self.compile(filename)
                    break # Inner loop
            else:
                break # Outer loop

    def getstatus(self, filename):
        """Return the compilation status for a file.

        The argument should be a filename normalized by
        os.path.realpath(), as returned by addfile().

        The return value is a string, with the following meanings:

        'todo'  -- not yet compiled
        'busy'  -- in the middle of compilation
        'done'  -- compiled successfully
        'error' -- compilation abandoned due to an error
        """
        assert filename == os.path.realpath(filename), filename
        assert filename.endswith(".fuel"), filename
        return self.filestatus.get(filename)

    # The remaining methods are not part of the public API,
    # except for check_entrypoint()

    def setstatus(self, filename, status):
        """Internal routine to set a file's compilation status.

        Call this with a normalized .fuel filename and a status value
        chosen from the list 'todo', 'busy', 'done', 'error'.  The
        file must have been added by addfile() first.  The following
        transitions are allowed: 'todo' -> 'busy'; 'busy' -> 'done';
        'busy' -> 'error'.  All other transitions are illegal.

        """
        assert filename == os.path.realpath(filename), filename
        assert filename.endswith(".fuel"), filename
        # Allow only transitions todo -> busy -> (done|error)
        oldstatus = self.filestatus.get(filename)
        if oldstatus == "todo":
            assert status == "busy", status
        else:
            assert oldstatus == "busy", oldstatus
            assert status in ("done", "error"), status
        self.filestatus[filename] = status

    def findmodule(self, name):
        """Override the base class method.

        This searches for .fuel files instead of .fuelc files,
        and normalizes the returned path.

        """
        filename =  modreg.ModuleRegistry.findmodule(self, name, [".fuel"])
        if filename is not None:
            filename = os.path.realpath(filename)
        return filename

    def loadmodule(self, name, filename):
        """Override the base class method.

        The filename argument should be a normalized .fuel file name.

        This method compiles the module unless the compiler output is
        already up to date and the force flag is off.

        """
        assert filename.endswith(".fuel")
        assert filename == os.path.realpath(filename), filename

        # Record the module name -> filename mapping
        if name in self.savedpaths:
            assert self.savedpaths[name] == filename
        else:
            self.savedpaths[name] = filename

        # Compute the package name from the module name
        i = name.rfind(".")
        if i >= 0:
            packagename = name[:i]
        else:
            packagename = None

        # There are several ways this can end:
        # 1) load a pre-existing .fuelc file (after checking dependencies)
        # 2) compile the .fuel file and load the resulting .fuelc file
        # 3) raise LoadError
        # The logic is a bit complicated, because there are several
        # different cases to start with:
        # a) (deleted)
        # b) .fuel file we haven't seen before (2)
        # c) .fuel file on our "to do" list (2)
        # d) .fuel file we've compiled before without errors (2)
        # e) .fuel file we've compiled before with errors (3)
        # f) .fuel file we're compiling already (recursive import) (3)

        # Check "todo" status
        fn_c = filename + "c"
        status = self.getstatus(filename)
        if status is None:
            # Case (b)
            # Never seen this file before; check .fuelc file if any,
            # except if self.force is set
            if not self.force:
                data, mtime = self.checkfiledeps(name, filename)
                if data is not None:
                    try:
                        return self.loadfrompickle(name, data)
                    except:
                        # If this fails, assume the pickle was bad
                        pass
            # We have to compile it first
            self.addfile(filename)
        else:
            if status == "done":
                # Case (d)
                # Previous compilation completed
                return self.loadfrompickle(name, self.loadpickle(fn_c))
            if status == "error":
                # Case (e)
                # Previous compilation failed
                raise modreg.LoadError("import %s: compilation error, "
                                       "please fix first" % name)
            if status == "busy":
                # Case (f)
                # Compilation still ongoing -- must be a recursive import
                raise modreg.LoadError("recursive import of %s" % name)
            # Case (c)
            assert status == "todo"

        # OK, let's compile it
        if not self.compile(filename, packagename):
            raise modreg.LoadError("import %s: compilation error, "
                                   "please fix first" % name)

        # No errors, so load it (yes, this reads what compile() just wrote!)
        return self.loadfrompickle(name, self.loadpickle(fn_c))

    def checkfiledeps(self, name, filename):
        """Check dependencies for a given file.

        The arguments are the module name and the filename.

        If all dependencies satisfied, this returns a tuple (data,
        mtime) where data is the unpickled data read from the .fuelc
        file and mtime is the mtime of the .fuelc file; otherwise this
        returns (None, None).

        """
        fuel_mtime = getmtime(filename)
        fn_c = filename + "c"
        fuelc_mtime = getmtime(fn_c)
        if fuelc_mtime >= fuel_mtime:
            # .fuelc file is up to date, check dependencies
            try:
                data = self.loadpickle(fn_c)
            except modreg.LoadError:
                pass # Oops, couldn't read .fuelc file or unpickle it
            else:
                if self.checkdeps(name, data, fuelc_mtime):
                    return data, fuelc_mtime
        return None, None

    def checkdeps(self, name, data, cutoff_mtime):
        """Check dependencies for a given module.

        The arguments are the module name and the unpickled data read
        from the .fuelc file.

        The return value is True if all dependencies are satisfied,
        False otherwise.

        """
        # Check dependencies.  data is a pickle read from *.fuel
        msg = logger.debug
        msg("checking dependencies for %s ...", name)
        dependencies = data.get("dependencies")
        if dependencies is None:
            msg("%s: no dependencies in .fuelc file", name)
            return False # An old .fuelc file, not yet containing dependencies
        assert isinstance(dependencies, dict)
        if ("fuel", "stdlib") not in dependencies and name != "fuel.stdlib":
            dependencies[("fuel", "stdlib")] = True
        for key, value in dependencies.items():
            msg("%s: %r -> %r", name, key, value)
            # There are two regular cases:
            # (1) key is a (prefix, name) tuple mapping to a bool;
            #     this corresponds to 'from prefix import name';
            #     value is True if name is a submodule of package
            #     prefix, False if it is some other object.
            # (2) key is a string mapping to a filename;
            #     this corresponds to 'import key'.
            # There is one backwards compatibility case:
            # (x) key is a tuple mapping to a string; this is an old-style
            #     dependency and the .fuelc file is definitely out of date
            if isinstance(key, tuple):
                # Case (1) or (x)
                if not isinstance(value, bool):
                    # Case (x)
                    msg("%s: old-style dependency found", name)
                    return False
                # Case (1)
                prefix, subname = key
                filename = self.findmodule(prefix + "." + subname)
                issubmodule = bool(filename)
                if issubmodule != value:
                    msg("%s: 'from %s import %s' interpretation changed",
                        name, prefix, subname)
                    return False
                if not issubmodule:
                    filename = self.findmodule(prefix)
                    if filename is None:
                        msg("%s: module %s disappeared?!", name, prefix)
                        return False
            else:
                # Case (2)
                filename = value
                assert filename is not None, key
                assert filename.endswith(".fuel"), (key, filename)
                if not self.update and filename in self.filestatus:
                    msg("%s: %r file status %r",
                        name, filename, self.filestatus[filename])
                    return False # It's compiled or to be compiled explicitly
                currentfilename = self.findmodule(key)
                if currentfilename != filename:
                    msg("%s: search for %r changed from %r to %r",
                        name, key, filename, currentfilename)
                    return False # The search result is different now
                data, mtime = self.checkfiledeps(key, filename)
                if data is None:
                    msg("%s: dependencies for %s failed", name, key)
                    return False
                assert mtime, repr(mtime)
                if mtime > cutoff_mtime:
                    msg("%s: %s is newer than dependent", name, filename)
                    return False
            fn_c = filename + "c"
            mtime = getmtime(fn_c)
            if not mtime:
                msg("%s: can't get mtime for %s", name, fn_c)
                return False
            if mtime > cutoff_mtime:
                msg("%s: %s is newer than dependent", name, fn_c)
                return False
        msg("%s: all dependencies satisified okay", name)
        return True # No excuse to fail :-)

    nesting_level = 0

    def compile(self, filename, packagename=None):
        """Compile a file and update its compilation status.

        Arguments are the filename and optionally the package name to
        which the file belongs (defaulting to the unnamed toplevel
        package).

        Return True iff parsing and compilation succeeded without
        errors; set status accordingly.

        """
        if packagename is None:
            packagename = self.getpackagename(filename)
        self.nesting_level += 1
        try:
            return self._compile(filename, packagename)
        finally:
            self.nesting_level -= 1

    def _compile(self, filename, packagename):
        """Internal version of compile().

        This only exists to put the setting of self.nesting_level in a
        separate place.

        """
        if self.update:
            basename = os.path.basename(filename)
            assert basename.endswith(".fuel")
            name = basename[:-5]
            if packagename:
                name = packagename + "." + name
            data, mtime = self.checkfiledeps(name, filename)
            if data is not None:
                logger.debug("%suptodate %s",
                             " "*(self.nesting_level-1), filename)
                self.setstatus(filename, "busy")
                self.setstatus(filename, "done")
                return True
        dependencies = {}
        oldcallback = self.getnotecallback()
        try:
            self.setnotecallback(dependencies.__setitem__)
            self.setstatus(filename, "busy")
            logger.info("%scompiling %s",
                        " "*(self.nesting_level-1), filename)
            rv = self.parse_and_translate(filename, packagename=packagename)
            if rv is None:
                self.setstatus(filename, "error")
                logger.error("error(s) compiling %s", filename)
                return False
            text, syms = rv
            text = PROLOGUE + text
            try:
                self.check_entrypoint(syms)
            except modreg.LoadError:
                text += REG_EPILOGUE
            else:
                text += MAIN_EPILOGUE
            self.write_fuelc(filename, text, syms, dependencies)
            self.setstatus(filename, "done")
            logger.debug("done compiling %s", filename)
            return True
        finally:
            self.setnotecallback(oldcallback)

    # Bunch of stuff to record dependencies

    def tryimportmodule(self, name):
        """Override the base class method.

        This calls notedepend() for successful loads.

        """
        defn = modreg.ModuleRegistry.tryimportmodule(self, name)
        if defn is not None:
            self.notedepend(name, self.savedpaths[name])
        return defn

    def importfrom(self, prefix, name):
        """Override the base class method.

        This calls notedepend() for successful loads.

        """
        defn = modreg.ModuleRegistry.importfrom(self, prefix, name)
        assert defn is not None
        self.notedepend((prefix, name), bool(self.lookup(prefix + "." + name)))
        return defn

    notecallback = None

    def getnotecallback(self):
        """Get the note callback."""
        return self.notecallback

    def setnotecallback(self, callback):
        """Set the note callback.

        This should be None or a function whose argument signature is
        (key, value), where either key is a module name and value a
        filename, or key is a (prefix, name) tuple and value is a bool
        indicating whether 'from prefix import name' refers to a
        module prefix/name.fuel (True) or to an object name in a
        module prefix.fuel (False).

        The note callback's return value is ignored.

        This should only be used by compile(); the callback is
        actually a dictionary's __setitem__ method.

        """
        self.notecallback = callback

    def notedepend(self, key, value):
        """Call the note callback, if not None."""
        if self.notecallback is not None:
            self.notecallback(key, value)

    # Stuff to perform the actual act of compilation

    bootstrapping = False # Flag set by scan_stdlib to avoid recursion

    def parse_and_translate(self, filename, basesyms=None, packagename=None):
        """Parse and translate a file.

        Error messages are written to the logger.

        If an error occurred, return None.  Otherwise, return a tuple
        (text, translator).

        """
        # Parse the file
        try:
            fileroot = driver.parse_file(filename)
        except IOError as err:
            logger.error("IOError: %s", err)
            return None
        except parse.ParseError as err:
            (msg, t, v, context) = (err.msg, err.type, err.value, err.context)
            try:
                lineno, offset = context
            except:
                lineno = context
                offset = None
                line = None
            else:
                line = linecache.getline(filename, lineno)
            logger.error('  File "%s", line %s', filename, lineno)
            if line:
                while line and line[0].isspace():
                    line = line[1:]
                    if offset:
                        offset -= 1
                logger.error("    %s", line.rstrip())
                if offset is not None:
                    logger.error("    %s", " "*offset + "^")
            else:
                if offset is not None:
                    logger.error("    offset %s (type=%r, value=%r)",
                                 offset, t, v)
                else:
                    logger.error("    (type=%r, value=%r)", t, v)
            logger.error(msg)
            return None

        if basesyms is None:
            if self.bootstrapping:
                basesyms = symtab.builtin_symbols
            else:
                basesyms = self.scan_stdlib()
                if basesyms is None:
                    return None

        # Translate the parse tree
        translator = translate.Translator(next=basesyms, moduleregistry=self,
                                          packagename=packagename, major=True)
        translator.errorhandler.reset()
        text = translator.translate(fileroot)
        if translator.errorhandler.errors:
            return None
        else:
            return text, translator

    def check_entrypoint(self, modobj, entrypoint="main", argtuples=[]):
        """Check that the module defines the entry point correctly.

        raise LoadError if it's not okay.

        Arguments:

        modobj     -- the module object to check
        entrypoint -- optional string for the entry point name, default 'main'
        argtuples  -- optional list of (name, type) tuples of
                      arguments to check, default empty

        Note: all entry points must have a first argument of type
        fuel.stdlib.Host; this must not be included in argtuples.

        """
        if self.hosttype is None:
            if not self.bootstrapping:
                self.scan_stdlib() # It's set as a side effect of this call
            if self.hosttype is None:
                raise modreg.LoadError("stdlib doesn't define class Host?!?!")
        defn = modobj.lookup(entrypoint)
        if defn is None:
            raise modreg.LoadError("entry point %s not defined" % entrypoint)
        if not isinstance(defn.type, typesys.Function):
            raise modreg.LoadError(
                "entry point %s defined but it is not callable (%s)" %
                (entrypoint, defn.type))
        argnames = defn.type.argnames
        argtypes = defn.type.argtypes
        resulttype = defn.type.resulttype
        ndefaults = defn.type.ndefaults
        if resulttype is not typesys.BasicVoid:
            raise modreg.LoadError("entry point %s() is a func, not proc" %
                                   entrypoint)
        if argtypes[0] != self.hosttype:
            raise modreg.LoadError(
                "entry point %s(): first arg is not a Host (%s)" %
                (entrypoint, argtypes[0]))
        if len(argtuples) > len(argtypes) - 1:
            raise modreg.LoadError(
                "entry point %s(): at most %d argument%s expected" %
                (entrypoint, len(argtypes) - 1, plural(len(argtypes)-1)))
        if len(argtuples) < len(argtypes) - 1 - ndefaults:
            raise modreg.LoadError(
                "entry point %s(): at least %d argument%s expected" %
                (entrypoint, len(argtypes) - 1 - ndefaults,
                 plural(len(argtypes) - 1 - ndefaults)))
        for i in range(len(argtuples)):
            aname, atype = argtuples[i]
            argname = argnames[i+1]
            argtype = argtypes[i+1]
            if argname != aname:
                raise modreg.LoadError(
                    "entry point %s(): "
                    "argument %d has name %s, but %s was given" %
                    (entrypoint, i+1, argname, aname))
            if str(argtype) != atype:
                raise modreg.LoadError(
                    "entry point %s(): "
                    "argument %d has type '%s', but '%s' was given" %
                    (entrypoint, i+1, argtype, atype))

    stdlibsyms = None
    hosttype = None

    def scan_stdlib(self):
        """Return the stdlib symbol table.

        As a side effect, set self.hosttype.

        """
        syms = self.stdlibsyms
        if syms is None:
            assert not self.bootstrapping
            try:
                self.bootstrapping = True
                defn = self.importmodule("fuel.stdlib")
                assert defn.type is typesys.BasicModule
                assert isinstance(defn.value, modreg.Module)
            finally:
                self.bootstrapping = False
            self.stdlibsyms = syms = symtab.SymbolTable(symtab.builtin_symbols)
            for name, subdefn in defn.value.symbols.items():
                syms.adddefn(name, subdefn)
            hostdefn = syms.lookup("Host")
            if hostdefn is not None and hostdefn.type is typesys.BasicType:
                self.hosttype = hostdefn.value
        return syms

    def write_fuelc(self, fn_fuel, text, syms, dependencies):
        """Write the compiler output to various files.

        For a Fuel source code file named X.fuel, the generated Python
        source code is written to X_fuel.py; the symbol table and
        dependencies are written to X.fuelc as a pickled dict with
        keys 'symbols', 'entrypoints' and 'dependencies'.  The public
        entry points and public symbols are computed here from the
        symbol table passed in; the dependencies are passed in.

        Public symbols are all toplevel symbols whose name doesn't
        start with underscore and whose type isn't 'module'.

        Public entry points are public symbols that are procedures
        whose first argument type is fuel.stdlib.Host.

        """
        # Compute various filenames
        assert fn_fuel.endswith(".fuel")
        fn = fn_fuel[:-5]
        fn_fuelc = fn + ".fuelc"
        fn_fuel_py = fn + "_fuel.py"
        # Write Python source code to <foo>_fuel.py
        logger.debug("writing source to %s", fn_fuel_py)
        new = fn_fuel_py + ".new"
        f = open(new, "w", encoding="utf-8")
        f.write(text)
        f.close()
        if os.path.exists(fn_fuel_py):
            os.remove(fn_fuel_py)
        os.rename(new, fn_fuel_py)
        # Write pickled symbols to <foo>.fuelc
        data = {"dependencies": dependencies}
        symbols = {}
        entrypoints = {}
        for name, defn in syms.names.items():
            type = defn.type
            if not name.startswith("_") and type is not typesys.BasicModule:
                symbols[name] = defn
                # entry points are procs with arg 1 of type Host
                if (isinstance(type, typesys.Function) and
                    type.resulttype is typesys.BasicVoid):
                    argtypes = type.argtypes
                    if argtypes and argtypes[0] == self.hosttype:
                        entrypoints[name] = (type.argnames, argtypes)
        data["symbols"] = symbols
        data["entrypoints"] = entrypoints
        logger.debug("writing symtab to %s", fn_fuelc)
        new = fn_fuelc + ".new"
        f = open(new, "wb")
        pickle.dump(data, f, 2)
        f.close()
        if os.path.exists(fn_fuelc):
            os.remove(fn_fuelc)
        os.rename(new, fn_fuelc)
