engine.py

'''
 Copyright (c) 2025 HERE Europe B.V.
 See the LICENSE file in the root of this project for license details.
'''

import types
from typing import overload

from flatdata.generator.tree.builder import build_ast, build_ast_from_file
from flatdata.generator.tree.nodes.trivial.namespace import Namespace
from flatdata.generator.tree.nodes.node import Node
from flatdata.generator.tree.syntax_tree import SyntaxTree

from .generators.cpp import CppGenerator
from .generators.dot import DotGenerator
from .generators.python import PythonGenerator
from .generators.rust import RustGenerator
from .generators.flatdata import FlatdataGenerator
from .generators import BaseGenerator


class Engine:
    """
    Flatdata Generator Engine.
    Implements code generation from the given flatdata schema.
    """

    _GENERATORS: dict[str, type[BaseGenerator]] = {
        "cpp": CppGenerator,
        "dot": DotGenerator,
        "py": PythonGenerator,
        "rust": RustGenerator,
        "flatdata" : FlatdataGenerator
    }

    @classmethod
    def available_generators(cls) -> list[str]:
        """
        Lists names of available code generators.
        """
        return list(cls._GENERATORS.keys())

    @classmethod
    def from_file(cls, path: str) -> 'Engine':
        """
        Create Engine from a schema file, resolving imports.
        :raises FlatdataSyntaxError
        """
        engine = cls.__new__(cls)
        engine.tree = build_ast_from_file(path)
        engine.schema = engine.tree.root_schema or ""
        return engine

    def __init__(self, schema: str) -> None:
        """
        Instantiates generator engine for a given schema.
        :raises FlatdataSyntaxError
        """
        self.schema = schema
        self.tree = build_ast(schema)

    def render(self, generator_name: str) -> str:
        """
        Render schema with a given generator
        :param generator_name:
        """
        generator = self._create_generator(generator_name)
        if generator is None:
            raise ValueError(
                "Generator %s not implemented. Available options: %s" %
                ( generator_name, self.available_generators() )
            )

        output_content = generator.render(self.tree)
        return str(output_content)

    @overload
    def render_python_module(self, module_name: str | None, archive_name: str, root_namespace: str | None = None) -> tuple[types.ModuleType, type]: ...
    @overload
    def render_python_module(self, *, archive_name: str, root_namespace: str | None = None) -> tuple[types.ModuleType, type]: ...
    @overload
    def render_python_module(self, module_name: str | None = None, archive_name: None = None, root_namespace: str | None = None) -> types.ModuleType: ...

    def render_python_module(self, module_name: str | None = None, archive_name: str | None = None, root_namespace: str | None = None) -> types.ModuleType | tuple[types.ModuleType, type]:
        """
        Render python module.
        :param module_name: Module name to use. If none, root namespace name is used.
        :param archive_name: Archive name to lookup,
            if specified, archive type is returned along with the model
        :param root_namespace: Root namespace to pick in case of multiple top level namespaces.
        """
        ns = self._find_root_namespace(self.tree, archive_name, root_namespace)
        module_code = self.render("py")
        module = types.ModuleType(module_name if module_name is not None else ns.name)
        #pylint: disable=exec-used
        exec(module_code, module.__dict__)
        if archive_name is None:
            return module

        name = ns.name + "_" + archive_name
        archive_type = getattr(module, name)
        return module, archive_type

    @classmethod
    def _create_generator(cls, name: str) -> BaseGenerator | None:
        generator_type = cls._GENERATORS.get(name, None)
        if generator_type is None:
            return None

        return generator_type()  # type: ignore[call-arg]  # dict values are concrete subclasses with zero-arg __init__

    @staticmethod
    def _find_root_namespace(tree: SyntaxTree, archive_name: str | None, root_namespace: str | None = None) -> Namespace:
        root_children = tree.root.children
        root_namespaces = [
            child for child in root_children
            if isinstance(child, Namespace) and "builtin" not in child.name
        ]
        if not root_namespaces:
            raise RuntimeError("No root namespace found.")
        elif root_namespace:
            for namespace in root_namespaces:
                if namespace.name == root_namespace:
                    return namespace
            raise RuntimeError("Invalid root namespace provided. Could not find root namespace in archive.")
        elif archive_name:
            for namespace in root_namespaces:
                archive_path = namespace.name + Node.PATH_SEPARATOR + archive_name
                archive_node = namespace.get(archive_path)
                if archive_node is not None:
                    return namespace
            raise RuntimeError("Archive direct parent namespace is not a root namespace.")
        elif len(root_namespaces) > 1:
            raise RuntimeError("Ambiguous root namespace. Could not find root archive.")

        return root_namespaces[0]