# Copyright 2026 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.
# ==============================================================================
"""Generates MJCF schema using dropdown directives for XMLreference.rst.

Generates XMLschema.rst by reading xml_native_reader.cc and producing
nested dropdown directives with list-table for attributes.
"""

import copy
import sys

import os

# Map symbols to icons:
# ! = required element, can appear only once -> star (prominent, required)
# ? = optional element, can appear only once -> dot (minimal, single)
# * = optional element, can appear many times -> None (no icon, most common)
# R = optional element, can appear many times recursively -> sync (recursive)
SYMBOL_TO_ICON = {
    '!': 'star',
    '?': 'dot',
    '*': None,
    'R': 'sync',
}

ELEMENT_ORDER = [
    'mujoco',
    'option',
    'compiler',
    'size',
    'statistic',
    'asset',
    'body',
    'deformable',
    'contact',
    'equality',
    'tendon',
    'actuator',
    'sensor',
    'keyframe',
    'visual',
    'default',
    'custom',
    'extension',
]

# Special display names for elements (when different from the element name)
ELEMENT_DISPLAY_NAME = {
    'body': '(world)body',
}


def generate_dropdown(
    element_name: str,
    symbol: str,
    link_name: str,
    attributes: list[str],
    links: list[str],
    level: int,
    is_top_level: bool = False,
) -> str:
  """Generate a dropdown directive for an element with its attributes."""
  indent = '   ' * level
  output = ''

  display_name = ELEMENT_DISPLAY_NAME.get(element_name, element_name)
  icon = SYMBOL_TO_ICON.get(symbol)
  # Element name is a :ref: link. Icon/macro goes on the right side.
  # Use |*| for * elements (even though it's empty) for future flexibility.
  element_link = f':ref:`{display_name}<{link_name}>`'
  if icon:
    title = f'{element_link} :octicon:`{icon}`'
  else:
    title = f'{element_link} |*|'
  output += f'{indent}.. dropdown:: {title}\n'
  if is_top_level:
    output += f'{indent}   :open:\n'
  output += '\n'

  content_indent = indent + '   '

  # Responsive grid for attributes (2-3-4-4: mobile-tablet-desktop-large)
  if attributes:
    output += f'{content_indent}.. grid:: 2 3 4 4\n'
    output += f'{content_indent}   :gutter: 0\n'
    output += '\n'

    for attr in attributes:
      att_link_name = f'{link_name}-{attr}'
      if att_link_name not in links:
        raise ValueError(
            f'Link for attribute {att_link_name} not found, update'
            ' XMLreference.rst'
        )
      output += f'{content_indent}   .. grid-item::\n'
      output += f'{content_indent}      :ref:`{attr}<{att_link_name}>`\n'
      output += '\n'

  output += '\n'
  return output


def generate() -> str:
  """Generates XMLschema.rst by parsing mjcf_table.inc.

  The schema is defined in mjcf_table.inc (generated from mjcf.schema) as a
  nested structure called MJCF[]. This function parses that structure and
  generates nested dropdown directives with list-tables for attributes.

  Returns:
    RST content with nested dropdown directives for the MJCF schema.
  """
  script_dir = os.path.dirname(os.path.abspath(__file__))
  repo_root = os.path.dirname(os.path.dirname(script_dir))
  filepath = os.path.join(repo_root, 'src', 'xml', 'generated', 'mjcf_table.inc')
  xmlfile = os.path.join(repo_root, 'doc', 'XMLreference.rst')

  # Collect all link targets from XMLreference.rst for validation.
  links = []
  with open(xmlfile, 'r', encoding='utf-8') as file:
    for line in file:
      if line.startswith('.. _'):
        links.append(line.strip()[4:-1])

  output = """..
  DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.

.. raw:: html

   <div class="schema-controls" style="margin-bottom: 1em;">
     <button onclick="document.querySelectorAll('.schema-small details').forEach(d => d.open = true)" class="sd-btn sd-btn-outline-primary sd-btn-sm">Expand All</button>
     <button onclick="document.querySelectorAll('.schema-small details details').forEach(d => d.open = false)" class="sd-btn sd-btn-outline-secondary sd-btn-sm" style="margin-left: 0.5em;">Collapse All</button>
   </div>

"""

  with open(filepath, 'r', encoding='utf-8') as file:
    to_strip = {' ', '{', '}', '"'}

    # Process each top-level element in a specific order for consistent output.
    for top_element in ELEMENT_ORDER:
      level = 0
      element: list[str] = []
      parent = ['', '', '', '', '']  # Track parent elements for link names.
      found_element = False
      file.seek(0)

      # Skip to the MJCF schema definition in the C++ source.
      for line in file:
        if 'std::vector<const char*> MJCF[] = {' in line.strip():
          break

      # Parse the schema structure.
      for line in file:
        if line.strip().endswith('};'):
          break  # End of schema definition.

        # Track nesting level using angle brackets in the C++ source.
        if '<' in line:
          level += 1
          continue
        if '>' in line:
          level -= 1
          continue

        # Parse element definition from the line.
        line_is_done = 1 if '}' in line else 0
        line_copy = copy.copy(line)
        for item in to_strip:
          line_copy = line_copy.replace(item, '')
        element += [item for item in line_copy.split(',') if item != '\n']

        # Determine if this element belongs to the current top-level element.
        outer_level = level == 0 and top_element == 'mujoco'
        top_level = level == 1 and element and top_element == element[0]
        sub_level = level > 1 and top_element == parent[2]
        level_is_correct = outer_level or top_level or sub_level

        if found_element and not level_is_correct:
          break  # Done with this top-level element.
        found_element = level_is_correct

        # Generate dropdown for completed element definition.
        if line_is_done and level_is_correct:
          # Link name: top-level uses element name, nested uses parent-child.
          link_name = element[0]
          if level > 1:
            link_name = parent[level] + '-' + element[0]
          if link_name not in links:
            raise ValueError(
                f'Link for element {link_name} not found, update'
                ' XMLreference.rst'
            )
          attributes = element[2:]  # First two items are name and symbol.

          # Adjust indentation level for RST output.
          dropdown_level = 0 if level == 0 else (1 if level == 1 else level)
          is_top_level = level == 0
          output += generate_dropdown(
              element[0],
              element[1],
              link_name,
              attributes,
              links,
              dropdown_level,
              is_top_level,
          )

        # Track parent elements for building nested link names.
        if line_is_done:
          if level < 4:
            parent[level + 1] = element[0]
          element = []

  # single newline at end of file
  return output.rstrip('\n') + '\n'


def main() -> int:
  if len(sys.argv) > 1:
    sys.exit('Too many command-line arguments.')
  sys.stdout.buffer.write(generate().encode('utf-8'))
  return 0


if __name__ == '__main__':
  main()