import shutil
import subprocess
import multiprocessing
import os
import sys

from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py as _build_py

base_dir = os.path.dirname(__file__)

with open(
    os.path.join(base_dir, "kuzu-source", "tools", "python_api", "requirements_dev.txt")
) as f:
    requirements = f.read().splitlines()


def _get_kuzu_version():
    cmake_file = os.path.join(base_dir, "kuzu-source", "CMakeLists.txt")
    with open(cmake_file) as f:
        for line in f:
            if line.startswith("project(Kuzu VERSION"):
                raw_version = line.split(" ")[2].strip()
                version_nums = raw_version.split(".")
                if len(version_nums) <= 3:
                    return raw_version
                else:
                    dev_suffix = version_nums[3]
                    version = ".".join(version_nums[:3])
                    version += ".dev%s" % dev_suffix
                    return version


kuzu_version = _get_kuzu_version()
print("The version of this build is %s" % kuzu_version)


class CMakeExtension(Extension):
    def __init__(self, name: str, sourcedir: str = "") -> None:
        super().__init__(name, sources=[])
        self.sourcedir = os.path.abspath(sourcedir)


class CMakeBuild(build_ext):
    def build_extension(self, ext: CMakeExtension) -> None:
        self.announce("Building native extension...")
        # Pass the platform architecture for arm64 to cmake for
        # cross-compilation.
        env_vars = os.environ.copy()
        python_version = ".".join(
            (str(sys.version_info.major), str(sys.version_info.minor))
        )
        self.announce("Python version is %s" % python_version)
        env_vars["PYBIND11_PYTHON_VERSION"] = python_version
        env_vars["PYTHON_EXECUTABLE"] = sys.executable

        if sys.platform == "darwin":
            archflags = os.getenv("ARCHFLAGS", "")

            if len(archflags) > 0:
                self.announce("The ARCHFLAGS is set to '%s'." % archflags)
                if archflags == "-arch arm64":
                    env_vars["CMAKE_OSX_ARCHITECTURES"] = "arm64"
                elif archflags == "-arch x86_64":
                    env_vars["CMAKE_OSX_ARCHITECTURES"] = "x86_64"
                else:
                    self.announce("The ARCHFLAGS is not valid and will be ignored.")
            else:
                self.announce("The ARCHFLAGS is not set.")

            deploy_target = os.getenv("MACOSX_DEPLOYMENT_TARGET", "")
            if len(deploy_target) > 0:
                self.announce("The deployment target is set to '%s'." % deploy_target)
                env_vars["CMAKE_OSX_DEPLOYMENT_TARGET"] = deploy_target

        build_dir = os.path.join(ext.sourcedir, "kuzu-source")

        # Clean the build directory.
        subprocess.run(["make", "clean"], check=True, cwd=build_dir)

        try:
            num_cores = int(os.environ["NUM_THREADS"])
            self.announce(
                "Using %d cores for building the native extension." % num_cores
            )
        except:
            self.announce("NUM_THREADS is not set. Using all available cores.")
            num_cores = multiprocessing.cpu_count()

        # Build the native extension.
        full_cmd = ["make", "python", "NUM_THREADS=%d" % num_cores]
        subprocess.run(full_cmd, cwd=build_dir, check=True, env=env_vars)
        self.announce("Done building native extension.")
        self.announce("Copying native extension...")
        dst = os.path.join(ext.sourcedir, ext.name)
        shutil.rmtree(dst, ignore_errors=True)
        shutil.copytree(
            os.path.join(build_dir, "tools", "python_api", "build", ext.name), dst
        )
        self.announce("Done copying native extension.")


class BuildExtFirst(_build_py):
    # Override the build_py command to build the extension first.
    def run(self):
        self.run_command("build_ext")
        return super().run()


setup(
    name="kuzu",
    version=kuzu_version,
    install_requires=[],
    ext_modules=[CMakeExtension(name="kuzu", sourcedir=base_dir)],
    description="An in-process property graph database management system built for query speed and scalability.",
    license="MIT",
    long_description=open(os.path.join(base_dir, "README.md"), "r").read(),
    long_description_content_type="text/markdown",
    packages=["kuzu"],
    zip_safe=True,
    include_package_data=True,
    cmdclass={
        "build_py": BuildExtFirst,
        "build_ext": CMakeBuild,
    },
)
