mujoco

# Copyright 2025 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from introspect import enums as introspect_enums
from introspect import functions as introspect_functions

from wasm.codegen.generators import common
from wasm.codegen.generators import constants
from wasm.codegen.generators import enums
from wasm.codegen.generators import functions
from wasm.codegen.generators import structs


class BindingBuilder:
  """Builds WASM bindings for MuJoCo."""

  def __init__(self):
    self.markers_and_content = []

  def set_enums(self):
    """Generates and sets the enum bindings."""
    self.markers_and_content += enums.generate(
        list(introspect_enums.ENUMS.values())
    )
    return self

  def set_structs(self):
    """Generates and sets the struct bindings."""
    self.markers_and_content += structs.generate(constants.STRUCTS_TO_BIND)
    return self

  def set_functions(self):
    """Generates and sets the function wrappers and bindings."""

    functions_to_bind = []
    for name, func in introspect_functions.FUNCTIONS.items():
      if not functions.is_excluded_function_name(name):
        functions_to_bind.append(func)

    self.markers_and_content += functions.generate(functions_to_bind)
    return self

  def build_string(self, template_path: str) -> str:
    """Applies the accumulated bindings to a template file and returns the output."""
    with open(template_path, "r") as f:
      content = f.readlines()

    for marker, lines in self.markers_and_content:
      content = common.replace_lines_containing_marker(
          content, marker, lines
      )

    text = "".join(content)
    text = text.replace(
        "// limitations under the License.\n\n",
        "// limitations under the License.\n\n// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.\n\n",
        1,
    )
    return text

  def build_file(self, template_path: str, generated_path: str):
    """Applies the accumulated bindings to a template file and writes the output."""
    text = self.build_string(template_path)
    common.write_to_file(generated_path, text)