"""
Given a collection of node metadata, ensure that all nodes appear in the [omnigraph.node_types] section of the
extension.toml used by the owning extension to register its information.

Invoke this script by passing the location of the .json file containing the metadata for all node types and the location
of the extension.toml file to which the metadata should be populated.

    python ensure_node_types_in_toml.py
           --nodeInfo $BUILD/exts/omni.my.extension/ogn/nodes.json
           --toml $SRC/extensions/omni.my.extension/config/extension.toml

The extension.toml file will be modified to included a generated section with this format:

    # === GENERATED BY ensure_node_types_in_toml.py -- DO NOT MODIFY ===
    [omnigraph.node_types]
    "omni.my.extension.MyNodeType" = 1
    "omni.my.extension.MyOtherNodeType" = 1
    # === END OF GENERATED CODE ===

or if the "--allData" flag is set then this more verbose format will be used:

    # === GENERATED BY ensure_node_types_in_toml.py -- DO NOT MODIFY ===
    [omnigraph.node_types."omni.my.extension.MyNodeType"]
    version = 1
    language = "C++"
    description = "This is my node"

    [omnigraph.node_types."omni.my.extension.MyOtherNodeType"]
    version = 1
    language = "C++"
    description = "This is my other node"
    # === END OF GENERATED CODE ===

Note that this script explicitly does not handle the case of multiple versions of the same node type in the same
extension as that is also not handled by OmniGraph proper.

You might also want to use an intermediate directory, which will create an explicit tag when the .toml is regenerated
so that you can safely handle the case of regeneration after a direct edit of the .toml file itself. This will ensure
that a user cannot accidentally delete a node definition from the automatically generated section;

    python ensure_node_types_in_toml.py
           --nodeInfo $BUILD/exts/omni.my.extension/ogn/nodes.json
           --toml $SRC/extensions/omni.my.extension/config/extension.toml
           --intermediate $TOP/_build/intermediate

Lastly, the support for the toml package is only available through repo_man in the build, not in the standard path.
You can pass in the root directory of the repo_man module if it is needed to find the toml package.

    python ensure_node_types_in_toml.py
           --nodeInfo $BUILD/exts/omni.my.extension/ogn/nodes.json
           --toml $SRC/extensions/omni.my.extension/config/extension.toml
           --intermediate $TOP/_build/intermediate
           --repoMan $TOP/_repo/deps/repo_man
"""
import argparse
import json
import logging
import os
import sys
from pathlib import Path

from node_generator.utils import WritableDir

# Create a logger and selectively turn on logging if the OGN debugging environment variable is set
logger = logging.getLogger("add_nodes_to_toml")
logging_handler = logging.StreamHandler(sys.stdout)
logging_handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
logger.addHandler(logging_handler)
logger.setLevel(logging.INFO if os.getenv("OGN_PARSE_DEBUG") else logging.WARN)


START_MARKER = "# === GENERATED BY ensure_node_types_in_toml.py -- DO NOT MODIFY ==="
"""Special marker text identifying the start of the generated node type identifier code"""

SECTION_NAME = "omnigraph.node_types"
"""Name of the generated section of the .toml file"""

END_MARKER = "# === END OF GENERATED CODE ==="
"""Special marker text identifying the end of the generated node type identifier code"""


# ======================================================================
class TomlModifier:
    """Encapsulates the reading, modifying, and writing of the .toml file with the node information

    Attributes:
        changes_made: True iff processing the .toml file resulted in changes to it
        toml: The toml module root - a member of the class because it may need to be imported from an alternate location

    Internal Attributes:
        __all_data: If True then the .toml will include language and description in addition to type name and version
        __existing_types: Dictionary of node_type_name:version_number for node types found in the original .toml file
        __node_info: Dictionary of node_type_name:node_type_info generated by the build for this extension
        __node_info_path: Path pointing to the generated file containing the node type information for this extension
        __tag_path: Path pointing to a file to use to tag the operation as complete (so that edits to the .toml can
        trigger regeneration)
        __toml_path: Path pointing to the .toml file to be modified
    """

    # --------------------------------------------------------------------------------------------------------------
    def __init__(self):
        """Set up the information required for the operations - do nothing just yet"""
        # Construct the parsing information. Run the script with "--help" to see the usage.
        parser = argparse.ArgumentParser(
            description="Ensure that the supplied .toml file contains metadata for all of the nodes defined"
            " in the extension",
            formatter_class=argparse.RawTextHelpFormatter,
        )
        parser.add_argument(
            "-a",
            "--allData",
            action="store_true",
            help="Dump all metadata to the .toml file instead of just the version and node type name",
        )
        parser.add_argument(
            "-in",
            "--intermediate",
            action=WritableDir,
            const=None,
            type=Path,
            metavar="INTERMEDIATE_DIRECTORY",
            help="Directory into which temporary build information is stored",
        )
        parser.add_argument(
            "-ni",
            "--nodeInfoFile",
            type=Path,
            metavar="NODE_INFO_FILE.json",
            help=".json file generated by the build that contains the registration information for all of the nodes",
        )
        parser.add_argument(
            "-rm",
            "--repoMan",
            type=Path,
            metavar="REPO_MAN_DIR",
            help="Path to the repo_man support directory",
        )
        parser.add_argument(
            "-t",
            "--tomlFile",
            type=Path,
            metavar="TOML_FILE",
            help=".toml file that will contain the node information",
        )
        parser.add_argument("-v", "--verbose", action="store_true", help="Output the steps the script is performing")
        args = parser.parse_args()

        # If the script steps are to be echoed enable the logger and dump the script arguments as a first step
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        logger.info("Processing the arguments")
        logger.info("    Args = %s", args)

        # Set up the internal attributes. All path files are resolved to absolute paths for convenience.
        self.__all_data = args.allData
        self.__toml_path = args.tomlFile.resolve()
        self.__node_info_path = args.nodeInfoFile.resolve()
        self.__tag_path = args.intermediate / "extension.toml.built" if args.intermediate is not None else None
        if self.__tag_path is not None:
            self.__tag_path = self.__tag_path.resolve()
        self.changes_made = False
        self.__in_team_city = False
        self.__existing_types = {}
        self.__node_info = {}

        # The toml package is installed as part of repo_man, not directly available in the build, so it may not be
        # available here and may have to be found through a location supplied by the script arguments.
        try:
            import toml

            self.toml = toml
        except ModuleNotFoundError:
            self.toml = None
        # There is some information to get from repoMan when running through a build so that we can successfully
        # determine when a failure is fatal and when it is just part of the normal generation process.
        if args.repoMan is not None:
            try:
                python_path = args.repoMan.resolve().as_posix()
                sys.path.append(python_path)
                import omni.repo.man

                if self.toml is None:
                    toml = omni.repo.man.get_toml_module()
                    self.toml = toml
                self.__in_team_city = omni.repo.man.is_running_in_teamcity()
            except (ModuleNotFoundError, AttributeError):
                # If repoMan is inaccessible then issue a warning but continue on to avoid spurious failures
                if self.toml is None:
                    logger.warning(
                        "toml module could not be found natively or at path '%s', parsing cannot happen", python_path
                    )
                else:
                    logger.warning(
                        "Not able to determine if running in TeamCity without module at '%s', assuming not.",
                        python_path,
                    )
        logger.info("   Team City run = %s", self.__in_team_city)

    # --------------------------------------------------------------------------------------------------------------
    @property
    def needs_generation(self) -> bool:
        """Returns True iff the extension.toml is older than either the generator script itself or the nodes.json file.
        This is done here after several unsuccessful attempts to get the build structure to recognize when the file
        needs rebuilding. At worst it means running this script when it isn't necessary, but then returning immediately
        after checking the file modification times so hopefully no big deal.
        """
        # If the toml cannot be parsed no generation can happen
        if self.toml is None:
            return False

        this_file = Path(__file__)
        # If the nodes.json does not exist then no generation is needed because there are no nodes
        if not self.__node_info_path.exists():
            logger.info("Skipping generation - no nodes.json file")
            return False

        # If the .toml file does not exist it definitely needs generation
        if not self.__toml_path.exists():
            logger.info("Forcing generation - no .toml file")
            return True

        # If the tag file does not exist but should then generation has never been done so it needs to be done now
        if self.__tag_path is not None and not self.__tag_path.exists():
            logger.info("Forcing generation - missing tag file")
            return True

        # All four files exist. Regeneration is only needed if the .tag file is not the newest one
        this_file_mtime = this_file.stat().st_mtime
        node_info_mtime = self.__node_info_path.stat().st_mtime
        toml_mtime = self.__toml_path.stat().st_mtime
        tag_mtime = self.__tag_path.stat().st_mtime if self.__tag_path is not None else 0
        if tag_mtime < toml_mtime:
            logger.info("Forcing generation - .toml is newer %s than the tag file %s", toml_mtime, tag_mtime)
            return True
        if tag_mtime < node_info_mtime:
            logger.info("Forcing generation - tag file is older than nodes.json")
            return True
        if tag_mtime < this_file_mtime:
            logger.info("Forcing generation - tag file is older than the generation script")
            return True

        # No good reason was found to regenerate, so don't
        logger.info("Skipping generation - the .toml file is up to date")
        return False

    # --------------------------------------------------------------------------------------------------------------
    def read_files(self):
        """Read in the contents of the .toml and .json files and parse them for modification"""
        logger.info("Reading the extension's .toml file")
        contents = self.toml.load(self.__toml_path)
        try:
            sections = SECTION_NAME.split(".")
            self.__existing_types = contents
            for section in sections:
                self.__existing_types = self.__existing_types[section]
        except KeyError:
            self.__existing_types = {}

        logger.info("Reading the extension's .json node information file")
        try:
            with open(self.__node_info_path, "r", encoding="utf-8") as json_fd:
                self.__node_info = json.load(json_fd)["nodes"]
        except (IOError, json.JSONDecodeError):
            self.__node_info = {}

    # --------------------------------------------------------------------------------------------------------------
    def __existing_version_matches(self, node_type_name: str, version: int) -> bool:
        """Returns True iff the node type name has a version number in the .toml file matching the one passed in"""
        if node_type_name not in self.__existing_types:
            return False
        node_type_info = self.__existing_types[node_type_name]
        # If the abbreviated version of the metadata was used the version number is all there is
        if isinstance(node_type_info, int):
            return version == node_type_info
        # Otherwise extract the version number from the metadata dictionary
        try:
            return version == node_type_info["version"]
        except KeyError:
            return False

    # --------------------------------------------------------------------------------------------------------------
    def add_nodes(self):
        """Ensure the nodes that were passed in are present in the .toml file"""
        logger.info("Ensuring the node types are in the file")

        for node_type_name, node_type_info in self.__node_info.items():
            version = int(node_type_info["version"])
            if not self.__existing_version_matches(node_type_name, version):
                self.changes_made = True
                if self.__all_data:
                    new_item = {
                        "version": version,
                        "language": node_type_info["language"],
                        "description": node_type_info["description"],
                    }
                else:
                    new_item = version
                logger.info("   Found an unregistered type - %s = %s", node_type_name, new_item)
                self.__existing_types[node_type_name] = new_item

    # --------------------------------------------------------------------------------------------------------------
    def write_file(self):
        """Write the new contents of the .toml file back to the original location.
        The toml library dump() method cannot be used here as it would lose information like comments and formatting
        so instead it narrows its focus to the [[omnigraph]] section, surrounding it with fixed markers so that it
        can be easily identified and replaced using a text-based edit.
        """
        logger.info("Writing the file")
        if self.__in_team_city:
            raise AttributeError(
                f"The file {self.__toml_path} was not up to date in the merge request. Rebuild and add it."
            )
        with open(self.__toml_path, "r", encoding="utf-8") as toml_fd:
            raw_contents = toml_fd.readlines()
        raw_line_count = len(raw_contents)

        # Convert the node type list to a .toml format
        logger.info("    Inserting new section %s", self.__existing_types)
        # Build the structure from the bottom up to ensure the .toml has the correct nesting
        section_dict = self.__existing_types
        sections = SECTION_NAME.split(".")
        sections.reverse()
        for section in sections:
            section_dict = {section: dict(sorted(section_dict.items()))}
        inserted_section = self.toml.dumps(section_dict)

        # Scan the file to see if/where the generated section currently resides
        in_section = False
        section_start_index = -1
        section_end_index = -1
        for line_index, line in enumerate(raw_contents):
            if in_section and line.rstrip() == END_MARKER:
                in_section = False
                section_end_index = line_index
            if line.rstrip() == START_MARKER:
                in_section = True
                section_start_index = line_index
        logger.info("    Existing section location was %s, %s", section_start_index, section_end_index)

        if section_start_index >= 0 and section_end_index == -1:
            raise ValueError(
                f"The .toml file '{self.__toml_path}' was illegal - it had a start marker but no end marker"
            )
        if section_start_index < 0:
            section_start_index = raw_line_count
            section_end_index = section_start_index

        # Write the modified contents with the new generated section
        try:
            with open(self.__toml_path, "w", encoding="utf-8") as toml_fd:
                toml_fd.writelines(raw_contents[0:section_start_index])
                # If inserting at the end of the file then insert a blank line for readability
                if section_start_index == raw_line_count:
                    toml_fd.write("\n")
                toml_fd.write(f"{START_MARKER}\n")
                toml_fd.write(inserted_section)
                toml_fd.write(f"{END_MARKER}\n")
                toml_fd.writelines(raw_contents[section_end_index + 1 : raw_line_count])
                toml_fd.flush()  # Required to ensure the mtime is earlier than the tag file's
        except IOError as error:
            raise IOError(f"Failed to write back the .toml file '{self.__toml_path}'") from error

    # --------------------------------------------------------------------------------------------------------------
    def touch_tag_file(self):
        """Forces update of the tag file mtime."""
        try:
            # Tag the conversion as being complete so that a build process can properly manage dependencies.
            # This has to happen last to avoid a false positive where the tag file is older than the .toml
            logger.info("Touching the tag file %s", self.__tag_path)
            if self.__tag_path is not None:
                with open(self.__tag_path, "w", newline="\n", encoding="utf-8") as tag_fd:
                    tag_fd.write("This file tags the last time its .toml file was processed with node metadata")
        except IOError as error:
            raise IOError(f"Failed to write back the tag file '{self.__tag_path}'") from error


# ==============================================================================================================
def main_update_extension_toml():
    """Walk through the steps required to parse, modify, and write the .toml file."""
    modifier = TomlModifier()
    if modifier.needs_generation:
        modifier.read_files()
        modifier.add_nodes()
        if modifier.changes_made:
            modifier.write_file()
        # The tag file needs updating even if changes were not made so that it doesn't repeatedly try and fail to
        # regenerate every time the script runs.
        modifier.touch_tag_file()


# ==============================================================================================================
if __name__ == "__main__":
    main_update_extension_toml()
