import asteval
import carb
import carb.settings
import numpy as np
import omni.graph.core as cg
import omni.kit.pipapi
import omni.stageupdate
from pxr import Sdf

__IMPORT_WHITELIST__ = ["torch", "carb"]


def do_import(module, as_name=""):
    if not as_name:
        as_name = module
    as_name = as_name[as_name.rfind(".") + 1 :]
    if not module.startswith("omni."):
        exec(  # noqa: PLW0122
            f"""
try:
   globals()["{as_name}"] = __import__("{module}")
except ImportError:
   print("{module} not found. attempting to install...")
   omni.kit.pipapi.install("{module}")
   globals()["{as_name}"] = __import__("{module}")
"""
        )
    return as_name, globals()[as_name]


def filter_import_list(modules):
    out = []
    for m in modules:
        if m not in __IMPORT_WHITELIST__:
            carb.log_warn(f"OgnExpression: import of {m} is not allowed")
        else:
            out.append(m)
    return out


__expr_cache__ = {}
__stage_subscription__ = None


# dummy class for attribute placeholder
class Attrs:
    def __init__(self):
        return


class ExprCache:
    def __init__(self, node):
        global __stage_subscription__
        self.node = node
        if not __stage_subscription__:
            stage_update = omni.stageupdate.get_stage_update_interface()
            __stage_subscription__ = stage_update.create_stage_update_node(
                "OgnExpression", None, OgnExpression.on_detach, None, None, None, OgnExpression.on_removed
            )
        self.symTable = asteval.make_symbol_table(use_numpy=False, np=np)
        self.symTable["inputs"] = Attrs()
        # self.symTable["outputs"] = Attrs()
        self.eval = asteval.Interpreter(symtable=self.symTable)

    @staticmethod
    def get(node):
        pp = Sdf.Path(node.get_prim_path())
        if pp not in __expr_cache__:
            return None
        return __expr_cache__[pp]

    @staticmethod
    def clear(prim):
        if prim in __expr_cache__:
            # print("clear cache",prim)
            c = __expr_cache__[prim]
            c._stage_subscription = None  # noqa: PLW0212
            del __expr_cache__[prim]

    @staticmethod
    def init(node):
        pp = Sdf.Path(node.get_prim_path())
        cache = ExprCache(node)
        __expr_cache__[pp] = cache
        return cache


class OgnExpression:
    """Implementation of the OmniGraph expression node"""

    @staticmethod
    def get_node_type():
        return "Expression"

    @staticmethod
    def get_upstream_attr_or_self(attr):
        """Returns the list of attributes upstream from the given one, or itself if there is nothing upstream"""
        attributes_upstream = attr.get_upstream_connections()
        return attributes_upstream if attributes_upstream else [attr]

    @staticmethod
    def compute(context, node):
        # if not carb.settings.get_settings().get("/persistent/omnigraph/allowExprNodeExe"):
        #     return True

        need_init = False
        _cache = ExprCache.get(node)
        if not _cache:
            # print("initCache", Sdf.Path(node.get_prim_path()))
            _cache = ExprCache.init(node)
            if node.get_attribute_exists("import"):
                # print("import", Sdf.Path(node.get_prim_path()))
                import_list = OgnExpression.get_upstream_attr_or_self(node.get_attribute("import"))
                import_list = [m.strip() for m in context.get_attr_value(import_list).split(",") if len(m.strip())]
                for m in import_list:
                    try:
                        as_name, module = do_import(m)
                        _cache.symTable[as_name] = module
                    except (ImportError, KeyError, NameError):
                        carb.log_error("OgnExpression: failed importing " + m)
                        continue
            need_init = True

        attr_names = [n.get_name() for n in node.get_attributes()]
        in_attr_names = [a for a in attr_names if a.startswith("inputs:")]
        for ia in in_attr_names:
            attr = OgnExpression.get_upstream_attr_or_self(node.get_attribute(ia))
            setattr(_cache.symTable["inputs"], ia.split(":")[-1], context.get_attr_value(attr))

        out_attrs = Attrs()
        _cache.symTable["outputs"] = out_attrs
        # for oa in out_attr_names:
        #     attr = OgnExpression.get_upstream_attr_or_self(node.get_attribute(oa))
        #     setattr(_cache.symTable["outputs"], oa.split(":")[-1], context.get_attr_value(attr))

        if need_init and node.get_attribute_exists("initExpr"):
            init_expr = OgnExpression.get_upstream_attr_or_self(node.get_attribute("initExpr"))
            init_expr = context.get_attr_value(init_expr)
            _cache.eval(init_expr)
            # exec(init_expr)

        expr = OgnExpression.get_upstream_attr_or_self(node.get_attribute("updateExpr"))
        expr = context.get_attr_value(expr)
        _cache.eval(expr)
        # exec(expr)

        # lazy update of output values
        cach_out_attrs_dict = out_attrs.__dict__
        out_attr_names = [a for a in attr_names if a.startswith("outputs:")]
        for out_attr_name in out_attr_names:
            attr_short_name = out_attr_name[8:]  # without "outputs:" prefix
            if attr_short_name in cach_out_attrs_dict:
                cache_out_attr_val = cach_out_attrs_dict[attr_short_name]
                attr = OgnExpression.get_upstream_attr_or_self(node.get_attribute(out_attr_name))
                # print("!", out_attr_name, len(cache_out_attr_val), cache_out_attr_val[0])
                context.set_attr_value(cache_out_attr_val, attr)

        return True

    @staticmethod
    def on_shutdown():
        global __stage_subscription__
        __stage_subscription__ = None

    @staticmethod
    def on_detach():
        global __expr_cache__, __stage_subscription__
        __expr_cache__ = {}
        __stage_subscription__ = None

    @staticmethod
    def on_removed(prim):
        ExprCache.clear(prim)


cg.register_node_type(OgnExpression, 1)
