#!/usr/bin/env python3

# Copyright 2016 gRPC authors.
#
# 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.

from __future__ import print_function

import errno
import filecmp
import glob
import os
import os.path
import pprint
import shutil
import subprocess
import sys
import traceback
import uuid

# the template for the content of protoc_lib_deps.py
DEPS_FILE_CONTENT = """
# Copyright 2017 gRPC authors.
#
# 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.

# AUTO-GENERATED BY make_grpcio_tools.py!
CC_FILES={cc_files}

PROTO_FILES={proto_files}

CC_INCLUDES={cc_includes}
PROTO_INCLUDE={proto_include}

{commit_hash_expr}
"""

# expose commit hash suffix and prefix for check_grpcio_tools.py
COMMIT_HASH_PREFIX = 'PROTOBUF_SUBMODULE_VERSION="'
COMMIT_HASH_SUFFIX = '"'

EXTERNAL_LINKS = [
    ('@com_google_absl//', 'third_party/abseil-cpp/'),
    ('@com_google_protobuf//', 'third_party/protobuf/'),
    ('@utf8_range//:', 'third_party/utf8_range/'),
]

PROTOBUF_PROTO_PREFIX = '@com_google_protobuf//src/'

# will be added to include path when building grpcio_tools
CC_INCLUDES = [
    os.path.join('third_party', 'abseil-cpp'),
    os.path.join('third_party', 'protobuf', 'src'),
    os.path.join('third_party', 'utf8_range'),
]

# include path for .proto files
PROTO_INCLUDE = os.path.join('third_party', 'protobuf', 'src')

# the target directory is relative to the grpcio_tools package root.
GRPCIO_TOOLS_ROOT_PREFIX = 'tools/distrib/python/grpcio_tools/'

# Pairs of (source, target) directories to copy
# from the grpc repo root to the grpcio_tools build root.
COPY_FILES_SOURCE_TARGET_PAIRS = [
    ('include', 'grpc_root/include'),
    ('src/compiler', 'grpc_root/src/compiler'),
    ('third_party/abseil-cpp/absl', 'third_party/abseil-cpp/absl'),
    ('third_party/protobuf/src', 'third_party/protobuf/src'),
    ('third_party/utf8_range', 'third_party/utf8_range')
]

# grpc repo root
GRPC_ROOT = os.path.abspath(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', '..'))

# the directory under which to probe for the current protobuf commit SHA
GRPC_PROTOBUF_SUBMODULE_ROOT = os.path.join(GRPC_ROOT, 'third_party',
                                            'protobuf')

# the file to generate
GRPC_PYTHON_PROTOC_LIB_DEPS = os.path.join(GRPC_ROOT, 'tools', 'distrib',
                                           'python', 'grpcio_tools',
                                           'protoc_lib_deps.py')

# the script to run for getting dependencies
BAZEL_DEPS = os.path.join(GRPC_ROOT, 'tools', 'distrib', 'python',
                          'bazel_deps.sh')

# the bazel target to scrape to get list of sources for the build
BAZEL_DEPS_PROTOC_LIB_QUERY = '@com_google_protobuf//:protoc_lib'

BAZEL_DEPS_COMMON_PROTOS_QUERIES = [
    '@com_google_protobuf//:well_known_type_protos',
    # has both plugin.proto and descriptor.proto
    '@com_google_protobuf//:compiler_plugin_proto',
]


def protobuf_submodule_commit_hash():
    """Gets the commit hash for the HEAD of the protobuf submodule currently
     checked out."""
    cwd = os.getcwd()
    os.chdir(GRPC_PROTOBUF_SUBMODULE_ROOT)
    output = subprocess.check_output(['git', 'rev-parse', 'HEAD'])
    os.chdir(cwd)
    return output.decode("ascii").splitlines()[0].strip()


def _bazel_query(query):
    """Runs 'bazel query' to collect source file info."""
    print('Running "bazel query %s"' % query)
    output = subprocess.check_output([BAZEL_DEPS, query])
    return output.decode("ascii").splitlines()


def _pretty_print_list(items):
    """Pretty print python list"""
    formatted = pprint.pformat(items, indent=4)
    # add newline after opening bracket (and fix indent of the next line)
    if formatted.startswith('['):
        formatted = formatted[0] + '\n ' + formatted[1:]
    # add newline before closing bracket
    if formatted.endswith(']'):
        formatted = formatted[:-1] + '\n' + formatted[-1]
    return formatted


def _bazel_name_to_file_path(name):
    """Transform bazel reference to source file name."""
    for link in EXTERNAL_LINKS:
        if name.startswith(link[0]):
            filepath = link[1] + name[len(link[0]):].replace(':', '/')

            # For some reason, the WKT sources (such as wrappers.pb.cc)
            # end up being reported by bazel as having an extra 'wkt/google/protobuf'
            # in path. Removing it makes the compilation pass.
            # TODO(jtattermusch) Get dir of this hack.
            return filepath.replace('wkt/google/protobuf/', '')
    return None


def _generate_deps_file_content():
    """Returns the data structure with dependencies of protoc as python code."""
    cc_files_output = _bazel_query(BAZEL_DEPS_PROTOC_LIB_QUERY)

    # Collect .cc files (that will be later included in the native extension build)
    cc_files = []
    for name in cc_files_output:
        if name.endswith('.cc'):
            filepath = _bazel_name_to_file_path(name)
            if filepath:
                cc_files.append(filepath)

    # Collect list of .proto files that will be bundled in the grpcio_tools package.
    raw_proto_files = []
    for target in BAZEL_DEPS_COMMON_PROTOS_QUERIES:
        raw_proto_files += _bazel_query(target)
    proto_files = [
        name[len(PROTOBUF_PROTO_PREFIX):].replace(':', '/')
        for name in raw_proto_files
        if name.endswith('.proto') and name.startswith(PROTOBUF_PROTO_PREFIX)
    ]

    commit_hash = protobuf_submodule_commit_hash()
    commit_hash_expr = COMMIT_HASH_PREFIX + commit_hash + COMMIT_HASH_SUFFIX

    deps_file_content = DEPS_FILE_CONTENT.format(
        cc_files=_pretty_print_list(sorted(cc_files)),
        proto_files=_pretty_print_list(sorted(set(proto_files))),
        cc_includes=_pretty_print_list(CC_INCLUDES),
        proto_include=repr(PROTO_INCLUDE),
        commit_hash_expr=commit_hash_expr)
    return deps_file_content


def _copy_source_tree(source, target):
    """Copies source directory to a given target directory."""
    print('Copying contents of %s to %s' % (source, target))
    # TODO(jtattermusch): It is unclear why this legacy code needs to copy
    # the source directory to the target via the following boilerplate.
    # Should this code be simplified?
    for source_dir, _, files in os.walk(source):
        target_dir = os.path.abspath(
            os.path.join(target, os.path.relpath(source_dir, source)))
        try:
            os.makedirs(target_dir)
        except OSError as error:
            if error.errno != errno.EEXIST:
                raise
        for relative_file in files:
            source_file = os.path.abspath(
                os.path.join(source_dir, relative_file))
            target_file = os.path.abspath(
                os.path.join(target_dir, relative_file))
            shutil.copyfile(source_file, target_file)


def main():
    os.chdir(GRPC_ROOT)

    # Step 1:
    # In order to be able to build the grpcio_tools package, we need the source code for the codegen plugins
    # and its dependencies to be available under the build root of the grpcio_tools package.
    # So we simply copy all the necessary files where the build will expect them to be.
    for source, target in COPY_FILES_SOURCE_TARGET_PAIRS:
        # convert the slashes in the relative path to platform-specific path dividers.
        # All paths are relative to GRPC_ROOT
        source_abs = os.path.join(GRPC_ROOT, os.path.join(*source.split('/')))
        # for targets, add grpcio_tools root prefix
        target = GRPCIO_TOOLS_ROOT_PREFIX + target
        target_abs = os.path.join(GRPC_ROOT, os.path.join(*target.split('/')))

        _copy_source_tree(source_abs, target_abs)
    print(
        'The necessary source files were copied under the grpcio_tools package root.'
    )
    print()

    # Step 2:
    # Extract build metadata from bazel build (by running "bazel query")
    # and populate the protoc_lib_deps.py file with python-readable data structure
    # that will be used by grpcio_tools's setup.py (so it knows how to configure
    # the native build for the codegen plugin)
    try:
        print('Invoking "bazel query" to gather the protobuf dependencies.')
        protoc_lib_deps_content = _generate_deps_file_content()
    except Exception as error:
        # We allow this script to succeed even if we couldn't get the dependencies,
        # as then we can assume that even without a successful bazel run the
        # dependencies currently in source control are 'good enough'.
        sys.stderr.write("Got non-fatal error:\n")
        traceback.print_exc(file=sys.stderr)
        return
    # If we successfully got the dependencies, truncate and rewrite the deps file.
    with open(GRPC_PYTHON_PROTOC_LIB_DEPS, 'w') as deps_file:
        deps_file.write(protoc_lib_deps_content)
    print('File "%s" updated.' % GRPC_PYTHON_PROTOC_LIB_DEPS)
    print('Done.')


if __name__ == '__main__':
    main()
