# 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 functions API for APIreference.rst."""

import re
import sys

import os
import sys
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_REPO_ROOT = os.path.dirname(os.path.dirname(_SCRIPT_DIR))
sys.path.insert(0, _SCRIPT_DIR)
import generate_api_header
_FUNCTIONS_OVERRIDE = 'doc/APIreference/functions_override.rst'

_RST_SECTION_REGEX = re.compile(r'^\.\. (?P<section>_[A-Z].+):')
_RST_FUNCTION_REGEX = re.compile(r'^\.\. _(?P<function>mj.+):')

rst_to_section = {
    '_Parseandcompile': ['Parse and compile', ''],
    '_Mainsimulation': ['Main simulation', ''],
    '_Support': ['Support', ''],
    '_Components': ['Components', ''],
    '_Subcomponents': ['Sub components', ''],
    '_Raycollisions': ['Ray casting', ''],
    '_Printing': ['Printing', ''],
    '_Virtualfilesystem': ['Virtual file system', ''],
    '_Assetcache': ['Asset cache', ''],
    '_Resources': ['Resources', ''],
    '_Initialization': ['Initialization', ''],
    '_Errorandmemory': ['Error and memory', ''],
    '_Miscellaneous': ['Miscellaneous', ''],
    '_Interaction': ['Interaction', ''],
    '_Visualization-api': ['Visualization', ''],
    '_OpenGLrendering': ['OpenGL rendering', ''],
    '_FilamentRenderingApi': ['Filament rendering', ''],
    '_UIframework': ['UI framework', ''],
    '_Derivatives-api': ['Derivatives', ''],
    '_Signeddistancefunction': ['Signed Distance Functions', ''],
    '_Plugins-api': ['Plugins', ''],
    '_Thread': ['Threads', ''],
    '_Standardmath': ['Standard math', ''],
    '_Vectormath': ['Vector math', ''],
    '_Sparsemath': ['Sparse math', ''],
    '_Quaternions': ['Quaternions', ''],
    '_Poses': ['Poses', ''],
    '_Decompositions': ['Decompositions / Solvers', ''],
    '_Attachment': ['Attachment', ''],
    '_AddTreeElements': ['Tree elements', ''],
    '_AddNonTreeElements': ['Non-tree elements', ''],
    '_Setactuatorparameters': ['Set actuator parameters', ''],
    '_AddAssets': ['Assets', ''],
    '_FindAndGetUtilities': ['Find and get utilities', ''],
    '_AttributeSetters': ['Attribute setters', ''],
    '_AttributeGetters': ['Attribute getters', ''],
    '_SpecUtilities': ['Spec utilities', ''],
    '_ElementInitialization': ['Element initialization', ''],
    '_ElementCasting': ['Element casting', ''],
}


def generate() -> str:
  """Generates functions.rst."""
  rst_str = ''
  api = generate_api_header.read_headers()

  filepath = os.path.join(_REPO_ROOT, _FUNCTIONS_OVERRIDE)
  with open(filepath, 'r', encoding='utf-8') as file:
    current_section = None
    current_function = None
    for line in file:

      # Try matching a section.
      match = _RST_SECTION_REGEX.search(line)
      if match is not None:
        current_section = match.group('section')
        if current_section in rst_to_section:
          current_function = None
        else:
          current_section = None
        continue

      # Try matching a function
      match = _RST_FUNCTION_REGEX.search(line)
      if match is not None:
        current_function = match.group('function')
        if current_function in api:
          api[current_function].doc = ''
          current_section = None
        else:
          current_function = None
        continue

      # Update section doc.
      if current_section is not None:
        rst_to_section[current_section][1] += line

      # Override doc.
      if current_function is not None:
        api[current_function].doc += line

  # Check that all function sections from the headers are mapped.
  mapped_sections = {v[0] for v in rst_to_section.values()}
  for token, defn in api.items():
    if defn.c_type == 'FUNCTION' and defn.section not in mapped_sections:
      raise ValueError(
          f'Function {token!r} is in unmapped section {defn.section!r}. '
          f'Add an entry to rst_to_section in generate_functions.py.'
      )

  # Write out RST file.
  rst_str += ("""
..
  AUTOGENERATED: DO NOT EDIT MANUALLY


""".lstrip())

  for section in rst_to_section:
    section_title = rst_to_section[section][0]
    rst_str += (
        f"""
.. {section}:

{section_title}
{'^'*len(section_title)}
{rst_to_section[section][1]}""".lstrip())

    for token in api:
      if api[token].c_type != 'FUNCTION' or api[token].section != section_title:
        continue

      doc = api[token].doc.strip()
      doc = doc.replace('&nbsp;', '``&nbsp;``')

      rst_str += (f"""
.. _{token}:

`{token} <#{token}>`__
{'~'*(2*len(token)+8)}

.. mujoco-include:: {token}

{doc}

""".lstrip())
  return rst_str


def main() -> None:
  if len(sys.argv) > 1:
    sys.exit('Too many command-line arguments.')

  sys.stdout.buffer.write(generate().encode('utf-8'))


if __name__ == '__main__':
  main()