import ast
import os
import sys
import types
import argparse
from pathlib import Path
from typing import Optional, Tuple
from types import ModuleType
import importlib
import importlib.util
from operator import attrgetter

from ._method_stubs import STUBS

TARGET_SUPERCLASSES = [
    "InboundAdapter",
    "BusinessService",
    "BusinessProcess",
    "BusinessOperation",
    "OutboundAdapter",
]

MESSAGE_SUPERCLASSES = [ "JsonSerialize", "PickleSerialize"]

# ——— local datatype map ———


DATATYPE_MAP = {"str": "%VarString", "int": "%Integer", "bool": "%Boolean"}

DATATYPE_MAP_Parameters = {"str": "STRING", "int": "INTEGER", "bool": "BOOLEAN"}


def snake_to_pascal(name: str) -> str:
    # Check if the string is in snake_case
    if "_" in name and (name.lower() == name or name.upper() == name):
        return "".join(word.capitalize() for word in name.split("_") if word)
    # Return original if not snake_case
    return name
    

def extract_iris_package_name(tree, default_name):
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == "iris_package_name":
                    if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
                        return node.value.value
    return default_name


def real_filename(path):
    """
    Given a path (possibly with the wrong case), return the Path
    object that reflects the actual on-disk filename.
    Works on case-insensitive and case-sensitive filesystems.
    """
    p = Path(path)
    parent = p.parent if p.parent.exists() else Path(".")
    try:
        # Iterate through actual entries in parent dir
        for entry in parent.iterdir():
            if entry.name.lower() == p.name.lower():
                # Found matching filename (correct casing on disk)
                return entry
    except OSError:
        pass
    # Fallback: just return what you passed in
    return p


def find_package_root_and_qualname(file_path: Path, source_root: Optional[Path] = None) -> Tuple[Path, str]:
    """
    Return (package_root, qualified_module_name) for a filesystem path to a .py file.

    If source_root is provided, the qualified name is computed relative to that root.
    Otherwise, we auto-detect by walking up directories as long as __init__.py exists.
    """
    p = Path(file_path).resolve()

    if source_root is not None:
        root = Path(source_root).resolve()
        rel = p.relative_to(root)
        parts = list(rel.parts)
        if parts[-1].endswith(".py"):
            parts[-1] = parts[-1][:-3]
        qualname = ".".join(parts)
        return root, qualname

    # Auto-detect root by climbing until no __init__.py
    cur = p.parent
    package_parts = []
    while (cur / "__init__.py").exists():
        package_parts.append(cur.name)
        cur = cur.parent
    package_parts.reverse()
    qualname = ".".join(package_parts + [p.stem]) if package_parts else p.stem
    return cur, qualname


def import_module_from_path(file_path: Path, qualname: str, package_root: Path) -> ModuleType:
    """
    Import a module from a file under the desired qualified name, ensuring __package__
    is correct so explicit relative imports (from .x import y) work.
    """
    if str(package_root) not in sys.path:
        sys.path.insert(0, str(package_root))

    spec = importlib.util.spec_from_file_location(qualname, str(file_path))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not create spec for {file_path} as {qualname}")

    module = importlib.util.module_from_spec(spec)
    # Ensure __package__ is the parent package (critical for relative imports!)
    module.__package__ = qualname.rpartition(".")[0]
    sys.modules[qualname] = module
    spec.loader.exec_module(module)
    return module



def load_to_iris(cls_string, cls_name):
    try:
        # required = ["IRISUSERNAME", "IRISPASSWORD", "IRISNAMESPACE", "IRISINSTALLDIR"]
        # ensure_env(required)
        # set_derived_env_vars()

        import iris

        print(f"\n\nLoading {cls_name} to IRIS...")
        iris_stream = iris._Stream.GlobalCharacter._New()
        iris_stream.Write(cls_string)
        iris._SYSTEM.OBJ.LoadStream(iris_stream, "ck")
    except Exception as e:
        print(
            f"Failed to load {cls_name} to iris. It is possible that an env variable was not set correctly. You can try to "
            "manually load the .cls files generated by this script using -o /path/to/output_folder"
        )
        print(
            "Here are some possible error messages and their meaning:"
            "No module named 'iris': You need to set the PYTHONPATH environment variable to point to /iris_install_dir/lib/python"
        )
        print(e)


# _______________________________________________________________________________________________________________________ load_to_iris #


def get_class_module_ast(cls) -> ast.Module:
    import inspect

    """Return the AST of the .py file where superclass is defined (assumes UTF-8)."""
    file_path = inspect.getsourcefile(cls)
    if file_path is None:
        raise RuntimeError(f"Cannot find a .py source file for {cls!r}")

    return Path(file_path).resolve()


# _______________________________________________________________________________________________________________________ get_class_module_ast #

def detect_custom_classes(tree, loaded_module: ModuleType, output, manual, real_path: Path,args_module):
    """
    Walk the AST and detect classes whose runtime superclasses are custom classes
    (determined through attributes on the runtime class). 
    """
    result = []
    superclass_path_list = set()
    superclass_module = {}

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            for base in node.bases:
                name = getattr(base, "id", getattr(base, "attr", None))
                if name in TARGET_SUPERCLASSES:
                    # handled by generate_os_classes
                    continue
                clsobj = getattr(loaded_module, node.name, None)
                if clsobj is not None:
                    hostname = getattr(clsobj, "_hostname", None)
                    if hostname in TARGET_SUPERCLASSES:
                        superclass = attrgetter(ast.unparse(base))(loaded_module)
                        superclass_name = superclass._fullname
                        superclass_path = get_class_module_ast(superclass)
                        if superclass_path not in _visited_paths:
                            superclass_path_list.add(superclass_path)
                            superclass_module[superclass_path] = str(superclass.__module__)

                        result.append((node.name, superclass_name, node, hostname))

    for superclass_path in superclass_path_list:
        if superclass_path in _visited_paths:
            continue
        next_argv = []
        if output:
            next_argv += ["-o", output]
        if manual:
            next_argv += ["--manual"]
        if args_module:
            next_argv += ["-m"]
            next_argv += [superclass_module[superclass_path]]
        else:
            next_argv += [str(superclass_path)]
        main(next_argv)

    return result


def generate_custom_classes(tree, script_name, folder_name, iris_package_name, output, script_path, manual, real_path, python_library, loaded_module: ModuleType,args_module):
    classes = detect_custom_classes(tree, loaded_module, output, manual, real_path,args_module)
    if not classes:
        return
    class_tmpl = STUBS.get("ClassDefinition")
    common_oninit = STUBS.get("Common", {}).get("OnInit", "")
    bp_oninit = STUBS.get("Common", {}).get("OnInitBP", "")

    all_classes = {}
    for cls_name, supercls, node, hostname in classes:
        props_lines, settings_list, message_map_methods = extract_props_and_settings(node, real_path)
        props_block = "\n".join(props_lines) if props_lines else ""
        params_settings = f'Parameter SETTINGS = "{",".join(settings_list)}";' if settings_list else ""
        param_lines = extract_params(node)
        param_lines.append(params_settings)
        params_block = "\n".join(param_lines) if param_lines else ""

        if hostname == "BusinessProcess":
            oninit = bp_oninit.format(
                ClassName=cls_name,
                PackageName=iris_package_name,
                PythonLibrary=python_library,
                Superclass=supercls,
                ScriptPath=script_path,
                ScriptName=script_name,
            )
        else:
            oninit = common_oninit.format(
                ClassName=cls_name,
                PackageName=iris_package_name,
                PythonLibrary=python_library,
                Superclass=supercls,
                ScriptPath=script_path,
                ScriptName=script_name,
            )

        methods = [oninit]
        super_stubs = STUBS.get(hostname, {})
        PascalName = ""

        for child in node.body:
            if isinstance(child, ast.FunctionDef):
                name = child.name
                args_string = get_args(child)
                if snake_to_pascal(name) in super_stubs:
                    stub_tmpl = super_stubs[snake_to_pascal(name)]
                elif (hostname == "BusinessOperation" and name not in message_map_methods):
                    continue
                elif (hostname in ("BusinessOperation", "OutboundAdapter") and "AnyMethod" in super_stubs):
                    PascalName = snake_to_pascal(name)
                    stub_tmpl = super_stubs["AnyMethod"]
                else:
                    continue

                content_stub = stub_tmpl.format(
                    ClassName=cls_name,
                    PackageName=iris_package_name,
                    PythonLibrary=python_library,
                    Superclass=supercls,
                    MethodName=name,
                    MethodNameModified = PascalName,
                    Arguments=args_string,
                    ScriptPath=script_path,
                    ScriptName=script_name,
                )
                methods.append(content_stub)

        methods_combined = "\n".join(methods)

        content = class_tmpl.format(
            ClassName=cls_name,
            PackageName=iris_package_name,
            Superclass=supercls,
            props=props_block,
            params=params_block,
            methods=methods_combined,
        )

        if manual and not output:
            script_path = Path(script_path).resolve()
            output = script_path / f"{iris_package_name}_classes"
            output.mkdir(parents=True, exist_ok=True)

        if output:
            out_path = os.path.join(output, f"{cls_name}.cls")
            with open(out_path, "w") as f:
                f.write(content)
            print(f"Generated {out_path}")

        all_classes[cls_name] = content

    if not manual:
        for cls_name in all_classes.keys():
            load_to_iris(all_classes[cls_name], cls_name)

# _______________________________________________________________________________________________________________________ generate_custom_classes #

def eval_node(node):
    try:
        return ast.literal_eval(node)
    except Exception:
        # Fallback if it's not a literal (e.g., a Name); keep a readable string
        return ast.unparse(node) if hasattr(ast, "unparse") else None
            

def get_args(func_node):
    args = func_node.args.args

    arg_list = []
    for arg in args:
        if arg.arg == "self":
            continue
        arg_list.append(arg.arg)

    arguments = ", ".join(arg_list)
    return arguments


def extract_params(node: ast.ClassDef):
    params = []

    for stmt in node.body:
        ann = None
        if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1:
            target_node = stmt.targets[0]
            call_node = stmt.value
        elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
            target_node = stmt.target
            call_node = stmt.value
            ann = stmt.annotation
        else:
            continue

        # now check it's really IRISParameter(...)
        if not (
            isinstance(call_node, ast.Call)
            and (
                (
                    isinstance(call_node.func, ast.Name)
                    and call_node.func.id == "IRISParameter"
                )
                or (
                    isinstance(call_node.func, ast.Attribute)
                    and call_node.func.attr == "IRISParameter"
                )
            )
        ):
            continue

        
        param_name = snake_to_pascal(target_node.id)


        from_args= {}
        IRISParameter_PARAM_ORDER = ["value", "datatype", "description", "keyword_list"]
        for i, arg in enumerate(call_node.args):
            if i < len(IRISParameter_PARAM_ORDER):
                from_args[IRISParameter_PARAM_ORDER[i]] = eval_node(arg)

        # pull literal kwargs out of the call
        kw = {a.arg: eval_node(a.value) for a in call_node.keywords}

        value = from_args.get("value",kw.get("value", None))
        dtype = from_args.get("datatype",kw.get("datatype", None))
        desc = from_args.get("description",kw.get("description", None))
        keyword_list = from_args.get("keyword_list",kw.get("keyword_list", None))
        
        if dtype is None and ann is not None:
            dtype = DATATYPE_MAP_Parameters.get(ast.unparse(ann))

        if desc:
            params.append(f"/// {desc}")

        # build the Property line
        line = f"Parameter {param_name} "
        if dtype:
            line += f" As {dtype} "
        if keyword_list is not None:
            line += f" [{keyword_list}]"
        if value is not None:
            if isinstance(value, str):
                line += f' = "{value}"'
            else:
                line += f" = {value}"
        else:
            try:
                value = ast.literal_eval(call_node.args[0])
                if isinstance(value, str):
                    line += f' = "{value}"'
                else:
                    line += f" = {value}"
            except:
                print("Failed to parse value from IRISParameter call for {node}")

        line += "; \n"
        params.append(line)

    return params


def message_map_xdata(MessageMap):
    output = []
    output.append("XData MessageMap")
    output.append("{")
    output.append("  <MapItems>")

    for msg_type, method in MessageMap.items():
        output.append(f'    <MapItem MessageType="{msg_type}">')
        output.append(f"      <Method>{snake_to_pascal(method)}</Method>")
        output.append("    </MapItem>")

    output.append("  </MapItems>")
    output.append("}")

    result = "\n".join(output)
    return result



def extract_props_and_settings(node: ast.ClassDef, real_path):

    props, settings, message_map_method = [], [], []

    for stmt in node.body:
        ann = None
        if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1:
            target_node = stmt.targets[0]
            call_node = stmt.value
        elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
            target_node = stmt.target
            call_node = stmt.value
            ann = stmt.annotation
        else:
            continue

        if target_node.id.lower() == "messagemap":
            message_map_as_dictionary = {
                k.value: v.value for k, v in zip(call_node.keys, call_node.values)
            }
            message_map_as_xdata_string = message_map_xdata(message_map_as_dictionary)
            props.append(message_map_as_xdata_string)
            for msg_type, method in message_map_as_dictionary.items():
                message_map_method.append(method)
            continue

        # now check it's really IRISProperty(...)
        if not (
            isinstance(call_node, ast.Call)
            and (
                (
                    isinstance(call_node.func, ast.Name)
                    and call_node.func.id == "IRISProperty"
                )
                or (
                    isinstance(call_node.func, ast.Attribute)
                    and call_node.func.attr == "IRISProperty"
                )
            )
        ):
            continue

        prop_name = snake_to_pascal(target_node.id)

        from_args= {}
        IRISProperty_PARAM_ORDER = [ "default", "datatype", "description", "settings"]
        for i, arg in enumerate(call_node.args):
            if i < len(IRISProperty_PARAM_ORDER):
                from_args[IRISProperty_PARAM_ORDER[i]] = eval_node(arg)

        # pull literal kwargs out of the call
        kw = {a.arg: eval_node(a.value) for a in call_node.keywords}

        default = from_args.get("default",kw.get("default", None))
        dtype = from_args.get("datatype",kw.get("datatype", None))
        desc = from_args.get("description",kw.get("description", None))
        prop_setting = from_args.get("settings",kw.get("settings", None))

        if prop_setting is not None:
            setting = ":".join(part.strip() for part in prop_setting.split(":"))
            if setting == "-":
                settings.append(f"-{prop_name}")
                continue # don't generate a Property in the backend as we intend to remove an inherited IRIS property from the UI
            if setting == "":
                settings.append(prop_name)
            else:
                settings.append(f"{prop_name}:{setting}")




        if dtype is None and ann is not None:
            dtype = ast.unparse(ann)

        # build the property description line
        if desc:
            props.append(f"/// {desc}")


        # build the Property line
        datatype = DATATYPE_MAP.get(dtype, '%VarString')
        line = f"Property {prop_name} As {datatype}"
        if default is not None:
            if datatype == "%VarString":
                line += f' [InitialExpression = "{default}"]'
            else:
                line += f" [InitialExpression = {default}]"

        line += "; \n"
        props.append(line)



    return props, settings, message_map_method


def find_ossubclasses(tree):
    result = []
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            for base in node.bases:
                name = getattr(base, "id", getattr(base, "attr", None))
                if name in TARGET_SUPERCLASSES:
                    result.append((node.name, name, node))
    return result


def generate_os_classes(tree, script_name, folder_name, iris_package_name, output, script_path, manual, real_path, python_library):
    classes = find_ossubclasses(tree)
    class_tmpl = STUBS.get("ClassDefinition")
    common_oninit = STUBS.get("Common", {}).get("OnInit", "")
    bp_oninit = STUBS.get("Common", {}).get("OnInitBP", "")

    all_classes = {}
    for cls_name, supercls, node in classes:
        props_lines, settings_list, message_map_methods = extract_props_and_settings(node, real_path)
        props_block = "\n".join(props_lines) if props_lines else ""
        params_settings = f'Parameter SETTINGS = "{",".join(settings_list)}";' if settings_list else ""
        param_lines = extract_params(node)
        param_lines.append(params_settings)
        params_block = "\n".join(param_lines) if param_lines else ""

        if supercls == "BusinessProcess":
            oninit = bp_oninit.format(
                ClassName=cls_name,
                PackageName=iris_package_name,
                PythonLibrary=python_library,
                Superclass=supercls,
                ScriptPath=script_path,
                ScriptName=script_name,
            )
        else:
            oninit = common_oninit.format(
                ClassName=cls_name,
                PackageName=iris_package_name,
                PythonLibrary=python_library,
                Superclass=supercls,
                ScriptPath=script_path,
                ScriptName=script_name,
            )

        methods = [oninit]
        super_stubs = STUBS.get(supercls, {})
        PascalName = ""

        for child in node.body:
            if isinstance(child, ast.FunctionDef):
                name = child.name
                args_string = get_args(child)
                if snake_to_pascal(name) in super_stubs: # this will take care of OnMessage and on_message in a business operation... 
                    stub_tmpl = super_stubs[snake_to_pascal(name)]
                elif (supercls == "BusinessOperation" and name not in message_map_methods):
                    continue
                elif (supercls in ("BusinessOperation", "OutboundAdapter") and "AnyMethod" in super_stubs):
                    PascalName = snake_to_pascal(name)
                    stub_tmpl = super_stubs["AnyMethod"]
                else:
                    continue

                content_stub = stub_tmpl.format(
                    ClassName=cls_name,
                    PackageName=iris_package_name,
                    PythonLibrary=python_library,
                    Superclass=supercls,
                    MethodName=name,
                    MethodNameModified = PascalName,
                    Arguments=args_string,
                    ScriptPath=script_path,
                    ScriptName=script_name,
                )
                methods.append(content_stub)

        methods_combined = "\n".join(methods)

        content = class_tmpl.format(
            ClassName=cls_name,
            PackageName=iris_package_name,
            Superclass="Ens." + supercls,
            props=props_block,
            params=params_block,
            methods=methods_combined,
        )

        if manual and not output:
            script_path = Path(script_path).resolve()
            output = script_path / f"{iris_package_name}_classes"
            output.mkdir(parents=True, exist_ok=True)

        if output:
            out_path = os.path.join(output, f"{cls_name}.cls")
            with open(out_path, "w") as f:
                f.write(content)
            print(f"Generated {out_path}")

        all_classes[cls_name] = content

    if not manual:
        for cls_name in all_classes.keys():
            load_to_iris(all_classes[cls_name], cls_name)


# _______________________________________________________________________________________________________________________ generate_os_classes #


def props_and_indices_from_msg_class(node: ast.ClassDef):
    props, indices = [], []

    for stmt in node.body:
        ann = None
        if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1:
            target_node = stmt.targets[0]
            call_node = stmt.value
        elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
            target_node = stmt.target
            call_node = stmt.value
            ann = stmt.annotation
        else:
            continue

        # now check it's really Column(...)
        if not (
            isinstance(call_node, ast.Call)
            and (
                (isinstance(call_node.func, ast.Name) and call_node.func.id == "Column")
                or (
                    isinstance(call_node.func, ast.Attribute)
                    and call_node.func.attr == "Column"
                )
            )
        ):
            continue

        prop_name = snake_to_pascal(target_node.id)


        from_args= {}
        Column_PARAM_ORDER = [ "default", "datatype", "description", "index"]
        for i, arg in enumerate(call_node.args):
            if i < len(Column_PARAM_ORDER):
                from_args[Column_PARAM_ORDER[i]] = eval_node(arg)

        # pull literal kwargs out of the call
        kw = {a.arg: eval_node(a.value) for a in call_node.keywords}

        default = from_args.get("default",kw.get("default", None))
        dtype = from_args.get("datatype",kw.get("datatype", None))
        desc = from_args.get("description",kw.get("description", None))
        index = from_args.get("index",kw.get("index", None))


        if dtype is None and ann is not None:
            dtype = ast.unparse(ann)

        # build the property description line
        if desc:
            props.append(f"/// {desc}")

        # build the Property line
        datatype = DATATYPE_MAP.get(dtype, "%VarString")
        line = f"Property {prop_name} As {datatype}"
        if default is not None:
            if datatype == "%VarString":
                line += f' [InitialExpression = "{default}"]'
            else:
                line += f" [InitialExpression = {default}]"

        line += "; \n"
        props.append(line)

        if index:
            index_line = f"INDEX {prop_name}IDX ON {prop_name}"
            index_line += "; \n"
            indices.append(index_line)

    return props, indices

def find_custom_message_classes(tree, loaded_module: ModuleType, output, manual, real_path: Path,args_module):
    result = []
    superclass_path_list = set()
    superclass_module = {}

    for node in ast.walk(tree):
        if not isinstance(node, ast.ClassDef):
            continue

        for base in node.bases:
            name = getattr(base, "id", None) or getattr(base, "attr", None)
            if name in MESSAGE_SUPERCLASSES:
                # handled by find_message_classes / generate_msg_wrappers
                continue

            clsobj = getattr(loaded_module, node.name, None)
            if clsobj is not None:
                serializer_class = getattr(clsobj, "_serializer_class", None)
                if serializer_class in MESSAGE_SUPERCLASSES:
                    superclass = attrgetter(ast.unparse(base))(loaded_module)
                    superclass_name = superclass._fullname
                    superclass_path = get_class_module_ast(superclass)
                    if superclass_path not in _visited_paths_msgs:
                        superclass_path_list.add(superclass_path)
                        superclass_module[superclass_path] = str(superclass.__module__)
                    result.append((node.name, superclass_name, node, serializer_class))

    for superclass_path in superclass_path_list:
        if superclass_path in _visited_paths_msgs:
            continue
        next_argv = []
        if output:
            next_argv += ["-o", output]
        if manual:
            next_argv += ["--manual"]
        if args_module:
            next_argv += ["-m"]
            next_argv += [superclass_module[superclass_path]]
        else:
            next_argv += [str(superclass_path)]
        main(next_argv)

    return result

def generate_custom_msg_wrappers(tree, script_name, folder_name, iris_package_name, python_library, output, script_path, manual, real_path, loaded_module: ModuleType,args_module):
    classes = find_custom_message_classes(tree, loaded_module, output, manual, real_path,args_module)
    if not classes:
        return

    all_classes = {}
    for cls, superclass_name, node, serializer in classes:
        props_lines, indices = props_and_indices_from_msg_class(node)
        props_block = "\n".join(props_lines) if props_lines else ""
        indices_block = "\n".join(indices) if indices else ""

        stub = STUBS.get("MsgSubclass")
        if not stub:
            print("Missing MsgWrapperSuperClass stub for " + superclass_name)
            continue
        content = stub.format(
            ClassName=cls,
            PackageName=iris_package_name,
            PythonLibrary=python_library,
            ScriptPath=script_path,
            ScriptName=script_name,
            props=props_block,
            indices=indices_block,
            Superclass=superclass_name
        )

        if manual and not output:
            script_path = Path(script_path).resolve()
            output = script_path / f"{iris_package_name}_classes"
            output.mkdir(parents=True, exist_ok=True)

        if output:
            out_path = os.path.join(output, f"{cls}.cls")
            with open(out_path, "w") as f:
                f.write(content)
            print(f"Generated {out_path}")

        all_classes[cls] = content

    if not manual:
        for cls_name in all_classes.keys():
            load_to_iris(all_classes[cls_name], cls_name)


def find_message_classes(source):
    """
    Parse the given Python source, walk the AST, and return a list of
    (class_name, matched_base) for every class that directly inherits from
    PickleSerialize, JsonSerialize.
    """
    tree = ast.parse(source)
    result = []

    for node in ast.walk(tree):
        if not isinstance(node, ast.ClassDef):
            continue

        # Check each base class for a match
        for base in node.bases:
            name = getattr(base, "id", None) or getattr(base, "attr", None)
            if name in MESSAGE_SUPERCLASSES:
                result.append((node.name, name, node))
                break  # stop on first match per class

    return result


def generate_msg_wrappers(tree, script_name, folder_name, iris_package_name, python_library, output, script_path, manual):
    classes = find_message_classes(tree)
    all_classes = {}
    for cls, superclass_name, node in classes:

        props_lines, indices = props_and_indices_from_msg_class(node)
        props_block = "\n".join(props_lines) if props_lines else ""
        indices_block = "\n".join(indices) if indices else ""

        stub = STUBS.get(superclass_name)
        if not stub:
            print("Missing MsgWrapperSuperClass stub for " + superclass_name)
            continue
        content = stub.format(
            ClassName=cls,
            PackageName=iris_package_name,
            PythonLibrary=python_library,
            ScriptPath=script_path,
            ScriptName=script_name,
            props=props_block,
            indices=indices_block,
        )

        if manual and not output:
            script_path = Path(script_path).resolve()
            output = script_path / f"{iris_package_name}_classes"
            output.mkdir(parents=True, exist_ok=True)

        if output:
            out_path = os.path.join(output, f"{cls}.cls")
            with open(out_path, "w") as f:
                f.write(content)
            print(f"Generated {out_path}")

        all_classes[cls] = content

    ## loading to IRIS. This might fail if env variables are not set correctly
    if not manual:
        for cls_name in all_classes.keys():
            load_to_iris(all_classes[cls_name], cls_name)


# _______________________________________________________________________________________________________________________ generate_msg_wrappers #



def get_package_root(module_name: str) -> Path:
    spec = importlib.util.find_spec(module_name)
    if not spec or not spec.origin or not spec.origin.endswith((".py", ".pyc")):
        raise ImportError(f"{module_name!r} does not resolve to a normal .py module (spec={spec})")

    base = Path(spec.origin).resolve().parent  # directory containing the .py file
    parts = module_name.split(".")
    project_root = base.parents[len(parts) - 2] if len(parts) > 1 else base
    return project_root

_visited_paths: set[Path] = set()
_visited_paths_msgs: set[Path] = set()


def main(argv: list[str] = None):
    parser = argparse.ArgumentParser(description="Generate ObjectScript classes")
    parser.add_argument("-o", "--output", required=False, help="Output folder")
    parser.add_argument("--manual", action="store_true", help="Run in manual mode")
    parser.add_argument("-m", "--module", help="Dotted module to analyze (e.g. pkg.sub.mod). If set, ignore positional file.")
    parser.add_argument("-s", "--source-root", help="Project source root used to compute absolute module names when loading from a file", dest="sourceroot")
    parser.add_argument("input_script", nargs="?", help="Path to a .py file (used when -m/--module is not provided)")
    args = parser.parse_args(argv)

    # Load source and module under a correct package context
    if args.sourceroot:
        sys.path.insert(0, os.path.abspath(args.sourceroot))
    else:
        # even though python automatically adds cwd to sys.path, running this cli tool DOES NOT.
        # So we need to manually add it. However, while recursively parsing imports, this gets fixed at
        # the iris namespace path (import iris does that)
        sys.path.insert(0, os.getcwd())

    if args.module:
        # Metadata for stubs
        module_name = args.module
        script_name = module_name
        script_path = get_package_root(module_name)
        # Module mode
        try:
            loaded_module = importlib.import_module(module_name)
        except Exception as e:
            print(f"Error importing module {module_name}: {e}")
            sys.exit(1)
        real_path = Path(loaded_module.__file__).resolve()
        try:
            source = real_path.read_text(encoding="utf-8")
        except Exception as e:
            print(f"Error reading source of module {module_name} at {real_path}: {e}")
            sys.exit(1)

    else:
        # File mode
        if not args.input_script:
            print("You must provide either -m/--module or a path to a .py file.")
            sys.exit(2)
        try:
            with open(args.input_script, "r", encoding="utf-8") as f:
                real_path = real_filename(f.name).resolve()
                source = f.read()
        except Exception as e:
            print(f"Error reading input script: {e}")
            sys.exit(1)

        # Compute package root + qualified name; import with that name
        package_root, module_name = find_package_root_and_qualname(real_path, Path(args.sourceroot).resolve() if args.sourceroot else None)
        try:
            loaded_module = import_module_from_path(real_path, module_name, package_root)
        except Exception as e:
            print(f"Error importing {real_path} as {module_name}: {e}")
            sys.exit(1)

        # Metadata for stubs
        script_name = real_path.stem
        script_path = real_path.parent

    # Parse AST
    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        print(f"Syntax error while parsing {real_path}: {e}")
        sys.exit(1)


    # Keep existing behaviour: if not inside package, fallback to script dir name or script name
    folder_name = os.path.basename(os.path.dirname(os.path.abspath(real_path))) or real_path.stem
    iris_package_name = extract_iris_package_name(tree, real_path.stem)
    python_library = module_name  # fully qualified module path

    if args.output:
        os.makedirs(args.output, exist_ok=True)

    # Track visited paths
    _visited_paths.add(real_path)
    _visited_paths_msgs.add(real_path)

    # Message wrappers (built-ins)
    generate_msg_wrappers(
        tree,
        script_name,
        folder_name,
        iris_package_name,
        python_library,
        args.output,
        script_path,
        args.manual,
    )

    # Message wrappers (custom)
    generate_custom_msg_wrappers(
        tree,
        script_name,
        folder_name,
        iris_package_name,
        python_library,
        args.output,
        script_path,
        args.manual,
        real_path,
        loaded_module,
        args.module
    )

    # Ens.* subclasses
    generate_os_classes(
        tree,
        script_name,
        folder_name,
        iris_package_name,
        args.output,
        script_path,
        args.manual,
        real_path,
        python_library,
    )

    # Custom subclasses of your runtime types
    generate_custom_classes(
        tree,
        script_name,
        folder_name,
        iris_package_name,
        args.output,
        script_path,
        args.manual,
        real_path,
        python_library,
        loaded_module,
        args.module
    )



if __name__ == "__main__":
    main()
