#!/usr/bin/env -S LD_PRELOAD=/usr/lib/libjemalloc.so /usr/bin/python3.13
#
# Copyright 2021 Pixar
#
# Licensed under the terms set forth in the LICENSE.txt file available at
# https://openusd.org/license.
#
"""
This script generates dynamic schema.usda, generatedSchema.usda and
plugInfo.json. The schema.usda is generated by parsing appropriate sdrNodes 
provided in a config file. Along with providing sdrNodes types and identifier, 
the config file also provides a list of subLayers which the auto populated 
schema.usda should sublayer. Code generation can also be optionally enabled via 
the json config. Note that code generation is disabled by default.

The script takes 3 arguments:
    - a json config, providing sourceType and sdrNodeIdentifiers
    - a destination directory. Note that a schema.usda with appropriate GLOBAL
      prim providing a libraryName in its customData, must be present at this
      location. This is also the location where generatedSchema.usda,
      plugInfo.json will get exported.
      In the case code generation is explicitly enabled in the config, 
      libraryPath must also be provided along with libraryName in the 
      customData, else usdGenSchema will fail.
    - an optional noreadme to ignore generating a README.md providing brief 
      explaination of the contents of the directory.

This script will run usdGenSchema on the auto populated schema.usda.
"""

from argparse import ArgumentParser, RawTextHelpFormatter
from textwrap import dedent
import os, sys, json
from subprocess import call
from pxr import Sdf, Tf, UsdUtils, Sdr
from pxr.UsdUtils.constantsGroup import ConstantsGroup
from pxr.UsdUtils.toolPaths import FindUsdBinary

class SchemaConfigConstants(ConstantsGroup):
    SDR_NODES = "sdrNodes"
    SOURCE_ASSET_NODES = "sourceAssetNodes"
    SUBLAYERS_STRING = "sublayers"
    RENDER_CONTEXT = "renderContext"

class SchemaLayerConstants(ConstantsGroup):
    GLOBAL_PRIM_PATH = "/GLOBAL"
    LIBRARY_NAME_STRING = "libraryName"
    SKIP_CODE_GENERATION = "skipCodeGeneration"
    USE_LITERAL_IDENTIFIER = "useLiteralIdentifier"
    SCHEMA_PATH_STRING = "schema.usda"

class MiscConstants(ConstantsGroup):
    USD_GEN_SCHEMA = "usdGenSchema"
    README_FILE_NAME = "README.md"

def _ConfigureSchemaLayer(schemaLayer, schemaSubLayers, skipCodeGeneration,
        useLiteralIdentifier):
    # - add sublayers
    subLayers = schemaLayer.subLayerPaths
    subLayersList = list(subLayers)
    subLayersList.extend(schemaSubLayers)
    schemaLayer.subLayerPaths = list(set(subLayersList))

    globalPrim = schemaLayer.GetPrimAtPath(
            SchemaLayerConstants.GLOBAL_PRIM_PATH)
    if not globalPrim:
        Tf.RaiseRuntimeError("Missing %s prim in schema.usda." \
                %(SchemaLayerConstants.GLOBAL_PRIM_PATH))

    if not globalPrim.customData:
        Tf.RaiseRuntimeError("customData must be defined on the %s prim" \
                %(SchemaLayerConstants.GLOBAL_PRIM_PATH))

    customDataDict = dict(globalPrim.customData)
    if SchemaLayerConstants.LIBRARY_NAME_STRING not in customDataDict:
        Tf.RaiseRuntimeError("customData on %s prim must provide a " \
            "%s." %(SchemaLayerConstants.GLOBAL_PRIM_PATH, 
                SchemaLayerConstants.LIBRARY_NAME_STRING))
    customDataDict[SchemaLayerConstants.SKIP_CODE_GENERATION] = \
            skipCodeGeneration
    customDataDict[SchemaLayerConstants.USE_LITERAL_IDENTIFIER] = \
            useLiteralIdentifier
            
    globalPrim.customData = customDataDict

    schemaLayer.Save()


if __name__ == '__main__':
    # Parse command line

    # config file provides:
    parser = ArgumentParser(description='''
    This script generates dynamic schema.usda, generatedSchema.usda and
    plugInfo.json. The schema.usda is generated by parsing appropriate sdrNodes 
    provided in a config file. Along with providing sdrNodes types and 
    identifier, the config file also provides a list of subLayers which the auto 
    populated schema.usda should sublayer. Code generation can also be 
    optionally enabled via the json config, note that code generation is 
    disabled by default.

    The script takes 3 arguments:
        - a json config, providing sdrNodes via sourceType, sdrNodeIdentifiers
          or explicit list of absolute asset file paths (sourceAssetNodes) and 
          a list of sublayers. The asset file paths might also contain
          environment variables, so users must set these prior to running the
          script.
        - a destination directory. Note that a schema.usda with appropriate 
          GLOBAL prim providing a libraryName in its customData, must be 
          present at this location. This is also the location where 
          generatedSchema.usda, plugInfo.json will get exported.
          In the case code generation is explicitly enabled in the config, 
          libraryPath must also be provided along with libraryName in the 
          customData, else usdGenSchema will fail.
        - an optional noreadme to ignore generating a README.md providing brief 
          explaination of the contents of the directory.

    This script will run usdGenSchema on the auto populated schema.usda.

    If regenerating schemas, it's recommended to set the
    USD_DISABLE_AUTO_APPLY_API_SCHEMAS environment variable to true in 
    order to prevent any previously generated auto-apply API schemas 
    from being applied to the specified schema bases which can result 
    in extra properties being pruned.

    The schema.usda populated specifications from the provided sdrNodes using
    UsdUtils.UpdateSchemaWithSdrNode and skipCodeGeneration metadata will be 
    set to true, unless explicitly marked False in the config for this 
    schema.usda.

    UsdUtils.UpdateSchemaWithSdrNode is responsible for:
    %s
    ''' %(UsdUtils.UpdateSchemaWithSdrNode.__doc__), 
    formatter_class=RawTextHelpFormatter)
    parser.add_argument('schemaConfig',
            nargs='?',
            type=str,
            default='./schemaConfig.json',
            help=dedent('''
            A json config providing sdrNodes via sourceType, sdrNodeIdentifiers
            or explicit list of absolute asset file paths (sourceAssetNodes).
            Note that for nodes specified under sourceAssetNodes we will use 
            the basename stripped of extension as the shaderId for nodes we 
            create. If node paths specified in sourceAssetNodes contain any
            environment variables, user is required to set these prior to 
            running the script.
            And also optionally providing a list of sublayers which the 
            schema.usda will sublayer. Code generation can also be optionally 
            enabled via the json config, note that code generation is disabled 
            by default. [Default: %(default)s]').
            Example json config file:
                {
                        "sdrNodes": 
                        {
                            "renderContext": "myRenderContext",
                            "sourceType": [
                                "sdrIdentifier1",
                                "sdrIdentifier2",
                                "sdrIdentifier3"
                            ],
                            "sourceAssetNodes": [
                                "/absolutepath/to/sdrNodeIdentifyingAsset1.extension,
                                "/absolutepath/to/sdrNodeIdentifyingAsset1.extension,
                            ],
                        },
                        "sublayers": [
                            "usd/schema.usda", 
                            "usdGeom/schema.usda", 
                            "usdLux/schema.usda"
                            ],
                        "skipCodeGeneration": True
                }
            '''))
    parser.add_argument('schemaGenerationPath',
            nargs='?',
            type=str,
            default='.',
            help=dedent('''
            The target directory where the code should be generated. The script
            assumes a basic schema.usda is defined at this location with a 
            GLOBAL prim configured with appropriate libraryName.
            [Default: %(default)s]
            '''))
    parser.add_argument('--noreadme',
            default=False,
            action='store_true',
            help=dedent('''
            When specified a README.md will not be created in the schemaGenerationPath
            explaining the source of the contents of this directory.
            '''))
    parser.add_argument('-v', '--validate',
            action='store_true',
            help=dedent('''
            This is passed to usdGenSchem to verify that the source files are 
            unchanged.
            '''))

    args = parser.parse_args()
    schemaGenerationPath = os.path.abspath(args.schemaGenerationPath)
    schemaConfig = os.path.abspath(args.schemaConfig)
    writeReadme = not args.noreadme
    validate = args.validate

    if not os.path.isfile(schemaConfig):
        Tf.RaiseRuntimeError("(%s) json config does not exist" %(schemaConfig))

    # Parse json config to extract sdrNodes and schema sublayers
    try:
        with open(schemaConfig) as json_file:
            config = json.load(json_file)
    except ValueError as ve:
        Tf.RaiseRuntimeError("Error loading (%s), value error: (%s)" \
                %(schemaConfig, ve))

    if not isinstance(config, dict):
        Tf.Warn("Invalid json config provided, expecting a dictionary")
        sys.exit(1)

    # Note that for nodes with explicit asset paths, this list stores a tuple,
    # with first entry being the sdrNode and second the true identifier
    # (identified from the file basename, which should match the node identifier
    # when queries using GetShaderNodeByIdentifierAndType at runtime).
    sdrNodesToParse = []
    renderContext = ""

    sdrNodesDict = config.get(SchemaConfigConstants.SDR_NODES)
    if sdrNodesDict:
        # Extract sdrNodes from the config
        sdrRegistry = Sdr.Registry()
        for sourceType in sdrNodesDict.keys():
            if sourceType == SchemaConfigConstants.RENDER_CONTEXT:
                # Extract any renderContext from the config if specified.
                renderContext = \
                    sdrNodesDict.get(SchemaConfigConstants.RENDER_CONTEXT)
            elif sourceType == SchemaConfigConstants.SOURCE_ASSET_NODES:
                # process sdrNodes provided by explicit sourceAssetNodes
                for assetPath in \
                sdrNodesDict.get(SchemaConfigConstants.SOURCE_ASSET_NODES):
                    assetPath = os.path.expandvars(assetPath)
                    node = Sdr.Registry().GetShaderNodeFromAsset(assetPath)
                    nodeIdentifier = \
                        os.path.splitext(os.path.basename(assetPath))[0]
                    if node:
                        sdrNodesToParse.append((node, nodeIdentifier))
                    else:
                        Tf.Warn("Node not found at path: %s." %(assetPath))
            else: 
                # we have an actual sdr node source type
                for nodeId in sdrNodesDict.get(sourceType):
                    node = sdrRegistry.GetShaderNodeByIdentifierAndType(nodeId,
                            sourceType)
                    if node is not None:
                        # This is a workaround to iterate through invalid 
                        # sdrNodes (nodes not having any input or output 
                        # properties). Currently these nodes return false when 
                        # queried for IsValid().
                        # Refer: pxr/usd/ndr/node.h#140-149
                        sdrNodesToParse.append(node)
                    else:
                        Tf.Warn("Invalid Node (%s:%s) provided." %(sourceType,
                            nodeId))
    else:
        Tf.Warn("No sdr nodes provided to generate a schema.usda")
        sys.exit(1)

    if not sdrNodesToParse:
        Tf.Warn("No valid sdr nodes provided to generate a schema.usda")
        sys.exit(1)

    schemaSubLayers = config.get(SchemaConfigConstants.SUBLAYERS_STRING, [])

    schemaLayerPath = os.path.join(schemaGenerationPath, 
            SchemaLayerConstants.SCHEMA_PATH_STRING)
    schemaLayer = Sdf.Layer.FindOrOpen(schemaLayerPath)
    if not schemaLayer:
        Tf.RaiseRuntimeError("Missing schema.usda at path (%s)." \
                %(schemaLayerPath))

    # set skipCodeGeneration  customData to true, unless explicitly marked
    # False in the config file
    skipCodeGeneration = config.get(
            SchemaLayerConstants.SKIP_CODE_GENERATION, True)
    # set useLiteralIdentifier customData to true, unless explicitly marked
    # False in the config file
    useLiteralIdentifier = config.get(
            SchemaLayerConstants.USE_LITERAL_IDENTIFIER, True)

    # configure schema.usda
    # fill in sublayers
    _ConfigureSchemaLayer(schemaLayer, schemaSubLayers, skipCodeGeneration,
            useLiteralIdentifier)

    # for each sdrNode call updateSchemaFromSdrNode with schema.usda
    for node in sdrNodesToParse:
        assetPathIdentifier = node[1] if type(node) is tuple else ""
        sdrNode = node[0] if type(node) is tuple else node
        UsdUtils.UpdateSchemaWithSdrNode(schemaLayer, sdrNode, renderContext,
                assetPathIdentifier)

    usdGenSchemaCmd = FindUsdBinary(MiscConstants.USD_GEN_SCHEMA)
    usdGenSchemaArgs = ["--validate"] if validate else []
    if not usdGenSchemaCmd:
        Tf.RaiseRuntimeError("%s not found. Make sure %s is in the PATH." \
                %(MiscConstants.USD_GEN_SCHEMA))

    call([usdGenSchemaCmd] + usdGenSchemaArgs, cwd=schemaGenerationPath)

    if writeReadme:
        readMeFile = os.path.join(schemaGenerationPath,
                MiscConstants.README_FILE_NAME)

        commonDescription = dedent("""
            The json config can provide sdrNodes either using sourceType and
            identifiers or using explicit paths via sourceAssetNodes. Note that
            if explicit paths contain any environment variables, then the user 
            is required to set these prior to running the script. Example:
            "$RMANTREE/lib/defaults/PRManAttribute.args", will require setting
            the RMANTREE environment variable before running the script.

            If regenerating schemas, it's recommended to set the
            USD_DISABLE_AUTO_APPLY_API_SCHEMAS environment variable to true in 
            order to prevent any previously generated auto-apply API schemas 
            from being applied to the specified schema bases which can result 
            in extra properties being pruned.

            Note that since users of this script have less control on direct
            authoring of schema.usda, "useLiteralIdentifier" is unconditionally
            set to true in schema.usda, which means the default camelCase token 
            names will be overriden and usdGenSchema will try keep the token 
            names as-is unless these are invalid.
            """)

        description = dedent("""
            The files ("schema.usda", "generatedSchema.usda" and
            "plugInfo.json") in this directory are auto generated using 
            usdgenschemafromsdr utility.

            A schema.usda is populated using sdrNodes which are specified in a
            json config. usdGenSchema is then run on this auto populated schema 
            (with skipCodeGeneration set to True) to output a 
            generatedSchema.usda and plugInfo.json.
            %s
            """)%(commonDescription) \
            if skipCodeGeneration else \
            dedent("""
            The files ("schema.usda", "generatedSchema.usda", "plugInfo.json",
            cpp source and header files) in this directory are auto generated
            using usdgenschemafromsdr utility.

            A schema.usda is populated using sdrNodes which are specified in a
            json config. usdGenSchema is then run on this auto populated schema 
            to output a generatedSchema.usda and plugInfo.json and all the
            generated code.
            %s
            """)%(commonDescription)

        with open(readMeFile, "w") as file:
            file.write(description)
    
