#!/bin/env python
#*****************************************************************************
# Copyright (c) 2017-2023, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#*****************************************************************************

# This script generates the MDL runtime header file for user modules.
#
# Call it like this:
# python gen_user_modules_runtime_header.py ../../compiler/stdmodule ../../../prod/bin/mdl_sdk_example/user_modules/mdl_runtime.h
#
# python 2.6 or higher is needed
#
import sys
import re
import os

from gen_intrinsic_func import SignatureParser, error

type_map = {
    "int2": "vint2",
    "int3": "vint3",
    "int4": "vint4",

    "float2": "vfloat2",
    "float3": "vfloat3",
    "float4": "vfloat4",

    "double2": "vdouble2",
    "double3": "vdouble3",
    "double4": "vdouble4",

    "bool2": "vbool2",
    "bool3": "vbool3",
    "bool4": "vbool4",
}

reference_parameter_types = {
    "bool2",
    "bool3",
    "bool4",
    "color",
    "double2",
    "double3",
    "double4",
    "float2",
    "float3",
    "float4",
    "int2",
    "int3",
    "int4"
}


def eat_until(token_set, tokens):
    """eat tokens until token_kind is found and return them, handle parenthesis"""
    r = 0
    e = 0
    g = 0
    a = 0
    l = len(tokens)
    eaten_tokens = []
    while l > 0:
        tok = tokens[0]
        if r == 0 and e == 0 and g == 0 and a == 0 and tok in token_set:
            return eaten_tokens, tokens
        if tok == '(':
            r += 1
        elif tok == ')':
            r -= 1
        elif tok == '[':
            e += 1
        elif tok == ']':
            e -= 1
        elif tok == '{':
            g += 1
        elif tok == '}':
            g -= 1
        elif tok == '[[':
            a += 1
        elif tok == ']]':
            a -= 1
        eaten_tokens.append(tokens[0])
        tokens = tokens[1:]
        l -= 1

    # do not return empty tokens, the parser do not like that
    return eaten_tokens, [None]


def format_param(param):
    typename, name, defparam = param

    typename = type_map.get(typename, typename)

    if typename in reference_parameter_types:
        res = "%s const &%s" % (typename, name)
    else:
        res = "%s %s" % (typename, name)

    if defparam:
        res += " = %s" % defparam
    return res


def parse_prototype(sigparser, decl, prototypes):
    """Get the C++ prototype for a given function declaration."""
    # poor man's scanner :-)
    tokens = re.sub(r'[,()]', lambda m: ' ' + m.group(0) + ' ', decl).split()

    tokens, ret_type = sigparser.get_type(tokens)

    name = tokens[0]

    if tokens[1] != '(':
        error("unknown token '" + tokens[1] + "' while processing '" + decl + "': '(' expected")
        sys.exit(1)

    tokens = tokens[2:]

    params = []

    if tokens[0] != ')':
        while True:
            tokens, t = sigparser.get_type(tokens)
            paramname = tokens[0]
            tokens = tokens[1:]

            if tokens[0] == '=':
                # default argument
                defarg, tokens = eat_until({',':None, ')':None}, tokens[1:])
            else:
                defarg = []

            params.append((t, paramname, ''.join(defarg)))

            if tokens[0] == ')':
                break
            if tokens[0] != ',':
                error("unknown token '" + tokens[1] + "' while processing '"
                        + decl + "': ',' expected")
                sys.exit(1)
            # skip the comma
            tokens = tokens[1:]

    ret_type = type_map.get(ret_type, ret_type)

    prototype = "%s %s(%s);" % (ret_type, name, ", ".join(map(format_param, params)))
    if "[" in ret_type or name == "transpose" or "[<N>]" in prototype or "color " in prototype:
        prototype = "// %s  (not supported yet)" % prototype
    prototypes.append(prototype)


def print_wrapped(parser, fileobj, line, wrap_pos = 99):
    """print the given line (provided without newline at end) and wrap it at wrap_pos,
       splitting the line at commas. Also handles commented out lines."""
    orig_line = line
    prefix = ""
    next_prefix = "//     " if line.startswith("//") else "    "

    while parser.indent * 4 + len(prefix) + len(line) >= wrap_pos:
        splitpos = line.rfind(',', 0, wrap_pos - parser.indent * 4 - len(prefix))
        if splitpos == -1:
            raise Exception("Unable to split line: %s" % orig_line)
        parser.write(fileobj, prefix + line[:splitpos + 1] + "\n")
        line = line[splitpos + 1:].lstrip()
        prefix = next_prefix

    parser.write(fileobj, prefix + line + "\n")


def usage(args):
    """print usage info and exit"""
    print "Usage: %s stdlib_directory outputfile" % args[0]
    return 1


def main(args):
    """Process one file and generate signatures."""
    if len(args) != 3:
        return usage(args)

    stdlib_dir = args[1]
    out_name   = args[2]
    strict     = True

    prototypes = []

    # monkey patch SignatureParser to generate signature names suitable for C++ header files
    SignatureParser.get_signature = (
        lambda self, decl: parse_prototype(self, decl, prototypes))

    try:
        parser = SignatureParser(args[0], stdlib_dir, out_name, strict)

        # copy the copyright from first 3 lines of state.cpp
        state_cpp_path = os.path.join(os.path.dirname(out_name), "state.cpp")
        with open(state_cpp_path) as f:
            copyright = "".join([next(f) for x in xrange(3)])

        with open(out_name, "w") as f:
            parser.write(f, copyright)
            parser.write(f, "\n#ifndef MDL_USER_MODULES_MDL_RUNTIME_H\n"
                "#define MDL_USER_MODULES_MDL_RUNTIME_H\n")

            for module_name in ["math", "debug"]:
                # clear list before parsing next module
                del prototypes[:]

                parser.parse(module_name)

                parser.write(f, "\nnamespace %s\n" % module_name)
                parser.write(f, "{\n")
                parser.indent += 1
                for prototype in prototypes:
                    print_wrapped(parser, f, prototype)
                parser.indent -= 1
                parser.write(f, "}\n")

            parser.write(f, "\n#endif  // MDL_USER_MODULES_MDL_RUNTIME_H\n")

    except Exception as e:
        error(str(e))
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))
