# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.

#
#   CXX=c++ python3 setup.py build develop
# or
#   CXX=c++ pip install . -vv
#
# Potentially also set TORCHBEAST_LIBS_PREFIX.

import os
import subprocess
import sys
import unittest

import numpy as np
import setuptools
from torch.utils import cpp_extension


PREFIX = os.getenv("CONDA_PREFIX")

if os.getenv("TORCHBEAST_LIBS_PREFIX"):
    PREFIX = os.getenv("TORCHBEAST_LIBS_PREFIX")
if not PREFIX:
    PREFIX = "/usr/local"


extra_compile_args = []
extra_link_args = []

protoc = f"{PREFIX}/bin/protoc"

grpc_objects = [
    f"{PREFIX}/lib/libgrpc++.a",
    f"{PREFIX}/lib/libgrpc.a",
    f"{PREFIX}/lib/libgpr.a",
    f"{PREFIX}/lib/libaddress_sorting.a",
]

include_dirs = cpp_extension.include_paths() + [np.get_include(), f"{PREFIX}/include"]
libraries = []

if sys.platform == "darwin":
    extra_compile_args += ["-stdlib=libc++", "-mmacosx-version-min=10.14"]
    extra_link_args += ["-stdlib=libc++", "-mmacosx-version-min=10.14"]

    # Relevant only when c-cares is not embedded in grpc, e.g. when
    # installing grpc via homebrew.
    libraries.append("cares")
elif sys.platform == "linux":
    libraries.append("z")

grpc_objects.append(f"{PREFIX}/lib/libprotobuf.a")


actorpool = cpp_extension.CppExtension(
    name="libtorchbeast.actorpool",
    sources=[
        "libtorchbeast/actorpool.cc",
        "libtorchbeast/rpcenv.pb.cc",
        "libtorchbeast/rpcenv.grpc.pb.cc",
    ],
    include_dirs=include_dirs,
    libraries=libraries,
    language="c++",
    extra_compile_args=["-std=c++17"] + extra_compile_args,
    extra_link_args=extra_link_args,
    extra_objects=grpc_objects,
)

rpcenv = cpp_extension.CppExtension(
    name="libtorchbeast.rpcenv",
    sources=[
        "libtorchbeast/rpcenv.cc",
        "libtorchbeast/rpcenv.pb.cc",
        "libtorchbeast/rpcenv.grpc.pb.cc",
    ],
    include_dirs=include_dirs,
    libraries=libraries,
    language="c++",
    extra_compile_args=["-std=c++17"] + extra_compile_args,
    extra_link_args=extra_link_args,
    extra_objects=grpc_objects,
)


def build_pb():
    # Hard-code rpcenv.proto for now.
    source = os.path.join(os.path.dirname(__file__), "libtorchbeast", "rpcenv.proto")
    output = source.replace(".proto", ".pb.cc")

    if os.path.exists(output) and (
        os.path.exists(source) and os.path.getmtime(source) < os.path.getmtime(output)
    ):
        return

    print("calling protoc")
    if (
        subprocess.call(
            [protoc, "--cpp_out=libtorchbeast", "-Ilibtorchbeast", "rpcenv.proto"]
        )
        != 0
    ):
        sys.exit(-1)
    if (
        subprocess.call(
            protoc + " --grpc_out=libtorchbeast -Ilibtorchbeast"
            " --plugin=protoc-gen-grpc=`which grpc_cpp_plugin`"
            " rpcenv.proto",
            shell=True,
        )
        != 0
    ):
        sys.exit(-1)


def test_suite():
    test_loader = unittest.TestLoader()
    test_suite = test_loader.discover("tests", pattern="*_test.py")
    return test_suite


class build_ext(cpp_extension.BuildExtension):
    def run(self):
        build_pb()

        cpp_extension.BuildExtension.run(self)


setuptools.setup(
    name="libtorchbeast",
    packages=["libtorchbeast"],
    version="0.0.13",
    author="TorchBeast team",
    ext_modules=[actorpool, rpcenv],
    cmdclass={"build_ext": build_ext},
    test_suite="setup.test_suite",
)
