import os
import shlex
import subprocess
import sys
from pathlib import Path


def _cmake_args_from_env():
    cmake_args = os.environ.get("CMAKE_ARGS")
    if not cmake_args:
        return []
    return shlex.split(cmake_args)


def _set_cmake_define(cmake_args, name, value):
    prefix = f"-D{name}="
    filtered = [arg for arg in cmake_args if not arg.startswith(prefix)]
    filtered.append(f"{prefix}{value}")
    return filtered


def _set_cmake_default_define(cmake_args, name, value):
    """Sets a default CMake define only when it is not already present in the arguments.

    Args:
        cmake_args: The existing CMake CLI arguments.
        name: The CMake define name.
        value: The default value to append.

    Returns:
        list[str]: A list of CMake arguments with the default define added only when absent.
    """
    prefix = f"-D{name}="
    if any(arg.startswith(prefix) for arg in cmake_args):
        return cmake_args
    return [*cmake_args, f"{prefix}{value}"]


def _default_parallel_jobs():
    """Returns default parallel jobs for CMake builds.

    Returns:
        int | None: None when CMAKE_BUILD_PARALLEL_LEVEL is set, else CPU count.
    """
    cmake_parallel = os.environ.get("CMAKE_BUILD_PARALLEL_LEVEL")
    if cmake_parallel:
        return None
    return os.cpu_count() or 1


def _detect_upstream_onnx_prefix():
    """Returns the install prefix of the pip-installed upstream onnx package, if any.

    This is purely informational — used to print a friendly message when the
    user passes ``--with-upstream-onnx``. The actual upstream onnx C++ library
    used by ``bench_load_file`` is fetched and built from source by CMake when
    ``-DONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX=ON`` (the pip-installed wheel does
    not ship ``ONNXConfig.cmake`` or a precompiled ``libonnx`` and therefore
    cannot satisfy ``find_package(ONNX)``).

    Returns:
        str | None: Directory containing the pip-installed onnx package, or
        ``None`` when the package is not importable.
    """
    try:
        import onnx  # noqa: PLC0415
    except ImportError:
        return None
    return os.path.dirname(onnx.__file__)


try:
    from setuptools import Command, Distribution, setup
except ModuleNotFoundError:
    try:
        from distutils.cmd import Command
        from distutils.core import Distribution, setup
    except ModuleNotFoundError:

        def _spawn(command, dry_run):
            """Prints and executes a command unless dry-run mode is enabled."""
            print(" ".join(shlex.quote(cmd_part) for cmd_part in command))
            if not dry_run:
                subprocess.run(command, check=True)

        def _run_build_ext_without_packaging(args):
            """Executes build_ext or build_benchmarks without setuptools or distutils support."""
            if not args or args[0] not in {"build_ext", "build_benchmarks"}:
                return False

            command = args[0]
            inplace = False
            cpp_tests = False
            dry_run = False
            build_temp = "build/temp"
            build_lib = "build/lib"
            gprof = False
            with_upstream_onnx = False
            no_kernels = False
            parallel = None

            i = 1
            while i < len(args):
                arg = args[i]
                if arg in {"--inplace", "-i"}:
                    inplace = True
                elif arg == "--cpp-tests":
                    cpp_tests = True
                elif arg in {"--dry-run", "-n"}:
                    dry_run = True
                elif arg.startswith("--build-temp="):
                    build_temp = arg.split("=", 1)[1]
                elif arg.startswith("--build-lib="):
                    build_lib = arg.split("=", 1)[1]
                elif arg == "--build-temp" and i + 1 < len(args):
                    build_temp = args[i + 1]
                    i += 1
                elif arg == "--build-lib" and i + 1 < len(args):
                    build_lib = args[i + 1]
                    i += 1
                elif arg == "--gprof":
                    gprof = True
                elif arg == "--with-upstream-onnx":
                    with_upstream_onnx = True
                elif arg == "--no-kernels":
                    no_kernels = True
                elif arg.startswith("--parallel="):
                    value = arg.split("=", 1)[1]
                    try:
                        parallel = int(value)
                    except ValueError:
                        raise ValueError(
                            f"Invalid value for --parallel: expected an integer, got {value!r}."
                        ) from None
                elif arg in {"--parallel", "-j"} and i + 1 < len(args):
                    value = args[i + 1]
                    try:
                        parallel = int(value)
                    except ValueError:
                        raise ValueError(
                            f"Invalid value for --parallel: expected an integer, got {value!r}."
                        ) from None
                    i += 1
                else:
                    raise ValueError(
                        f"Unsupported argument for {command}: {arg!r}. "
                        "Supported arguments include: --inplace, --dry-run, "
                        "--build-temp, --build-lib, --cpp-tests, "
                        "--gprof, --with-upstream-onnx, --no-kernels, --parallel."
                    )
                i += 1

            root = Path(__file__).resolve().parent
            build_temp_path = Path(build_temp).resolve()
            build_temp_path.mkdir(parents=True, exist_ok=True)
            if parallel is None:
                parallel = _default_parallel_jobs()

            if command == "build_benchmarks":
                print("running build_benchmarks")
                cmake_args = [
                    "cmake",
                    "-S",
                    str(root),
                    "-B",
                    str(build_temp_path),
                    "-DCMAKE_BUILD_TYPE=RelWithDebInfo",
                    "-DONNX_LIGHT_BUILD_BENCHMARKS=ON",
                    "-DONNX_LIGHT_BUILD_PYTHON=OFF",
                ]
                if gprof:
                    cmake_args.append("-DONNX_LIGHT_BENCH_GPROF=ON")
                if with_upstream_onnx:
                    onnx_prefix = _detect_upstream_onnx_prefix()
                    if onnx_prefix is not None:
                        print(f"--with-upstream-onnx: pip onnx detected at {onnx_prefix}")
                    print("--with-upstream-onnx: CMake will FetchContent+build upstream onnx")
                    cmake_args.append("-DONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX=ON")
                _spawn(cmake_args, dry_run)
                build_cmd = [
                    "cmake",
                    "--build",
                    str(build_temp_path),
                    "--target",
                    "bench_parse_serialize",
                ]
                if parallel is not None:
                    build_cmd += ["--parallel", str(parallel)]
                _spawn(build_cmd, dry_run)
                bench_bin = build_temp_path / "bench_parse_serialize"
                if not dry_run and bench_bin.exists():
                    print(f"\nBenchmark binary: {bench_bin}")
                    print(f"Usage: {bench_bin} -n 20 -t 1\n")
            else:
                print("running build_ext")
                install_prefix = root if inplace else Path(build_lib).resolve()
                cmake_args = _cmake_args_from_env()
                cmake_args = _set_cmake_default_define(cmake_args, "CMAKE_BUILD_TYPE", "Release")
                if cpp_tests:
                    cmake_args = _set_cmake_define(cmake_args, "ONNX_LIGHT_BUILD_TESTS", "ON")
                if no_kernels:
                    cmake_args = _set_cmake_define(cmake_args, "ONNX_LIGHT_BUILD_KERNELS", "OFF")
                if with_upstream_onnx:
                    onnx_prefix = _detect_upstream_onnx_prefix()
                    if onnx_prefix is not None:
                        print(f"--with-upstream-onnx: pip onnx detected at {onnx_prefix}")
                    print("--with-upstream-onnx: CMake will FetchContent+build upstream onnx")
                    cmake_args = _set_cmake_define(
                        cmake_args, "ONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX", "ON"
                    )
                else:
                    cmake_args = _set_cmake_define(
                        cmake_args, "ONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX", "OFF"
                    )
                _spawn(
                    [
                        "cmake",
                        "-S",
                        str(root),
                        "-B",
                        str(build_temp_path),
                        f"-DPython_EXECUTABLE={sys.executable}",
                        *cmake_args,
                    ],
                    dry_run,
                )
                build_cmd = ["cmake", "--build", str(build_temp_path), "--config", "Release"]
                if parallel is not None:
                    build_cmd += ["--parallel", str(parallel)]
                install_cmd = [
                    "cmake",
                    "--install",
                    str(build_temp_path),
                    "--prefix",
                    str(install_prefix),
                ]
                if cpp_tests:
                    # Build every library, extension and C++ test in one CMake
                    # invocation so progress is reported against one global total.
                    _spawn(build_cmd, dry_run)
                    _spawn(install_cmd, dry_run)
                    ctest_cmd = [
                        "ctest",
                        "--test-dir",
                        str(build_temp_path),
                        "--build-config",
                        "Release",
                        "--output-on-failure",
                    ]
                    if parallel is not None:
                        ctest_cmd += ["--parallel", str(parallel)]
                    _spawn(ctest_cmd, dry_run)
                else:
                    _spawn(build_cmd, dry_run)
                    _spawn(install_cmd, dry_run)
            return True

        if _run_build_ext_without_packaging(sys.argv[1:]):
            raise SystemExit(0) from None
        raise


class NoConfigDistribution(Distribution):
    """Skips setup.cfg and pyproject.toml parsing for setup.py commands."""

    def parse_config_files(self, _filenames=None):
        """Skips setuptools configuration file parsing."""
        return None


class BuildExt(Command):
    """Builds the extension with CMake."""

    description = "builds C++ extension with CMake"
    user_options = [
        ("inplace", "i", "build extension in the source tree"),
        ("build-temp=", "t", "temporary build directory"),
        ("build-lib=", "b", "build directory for platform-specific files"),
        ("cpp-tests", None, "build the C++ unit tests and run them with ctest after installing"),
        (
            "with-upstream-onnx",
            None,
            (
                "auto-detect the pip-installed onnx package prefix and forward it as "
                "CMAKE_PREFIX_PATH so find_package(ONNX) succeeds (used by C++ unit "
                "tests and benchmarks that compare against the upstream onnx library)"
            ),
        ),
        (
            "no-kernels",
            None,
            (
                "build the reduced extension set without the operator-kernel runtime "
                "and backend-test registries (sets -DONNX_LIGHT_BUILD_KERNELS=OFF); "
                "the reference runtime and backend-test helpers then raise an explicit "
                "error when used"
            ),
        ),
        ("parallel=", "j", "number of parallel build jobs"),
    ]
    boolean_options = ["inplace", "cpp-tests", "with-upstream-onnx", "no-kernels"]

    def initialize_options(self):
        """Initializes default values for command options."""
        self.inplace = False
        self.build_temp = None
        self.build_lib = None
        self.cpp_tests = False
        self.with_upstream_onnx = False
        self.no_kernels = False
        self.parallel = _default_parallel_jobs()

    def finalize_options(self):
        """Finalizes build directory paths for unspecified options."""
        build_base = "build"
        if self.build_temp is None:
            self.build_temp = os.path.join(build_base, "temp")
        if self.build_lib is None:
            self.build_lib = os.path.join(build_base, "lib")

    def run(self):
        """Runs CMake configure, build, and install commands."""
        root = Path(__file__).resolve().parent
        build_temp = Path(self.build_temp).resolve()
        build_temp.mkdir(parents=True, exist_ok=True)

        install_prefix = root if self.inplace else Path(self.build_lib).resolve()
        cmake_args = _cmake_args_from_env()
        cmake_args = _set_cmake_default_define(cmake_args, "CMAKE_BUILD_TYPE", "Release")
        if self.cpp_tests:
            cmake_args = _set_cmake_define(cmake_args, "ONNX_LIGHT_BUILD_TESTS", "ON")
        if self.no_kernels:
            cmake_args = _set_cmake_define(cmake_args, "ONNX_LIGHT_BUILD_KERNELS", "OFF")
        if self.with_upstream_onnx:
            onnx_prefix = _detect_upstream_onnx_prefix()
            if onnx_prefix is not None:
                print(f"--with-upstream-onnx: pip onnx detected at {onnx_prefix}")
            print("--with-upstream-onnx: CMake will FetchContent+build upstream onnx")
            cmake_args = _set_cmake_define(
                cmake_args, "ONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX", "ON"
            )
        else:
            cmake_args = _set_cmake_define(
                cmake_args, "ONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX", "OFF"
            )

        self.spawn(
            [
                "cmake",
                "-S",
                str(root),
                "-B",
                str(build_temp),
                f"-DPython_EXECUTABLE={sys.executable}",
                *cmake_args,
            ]
        )
        build_cmd = ["cmake", "--build", str(build_temp), "--config", "Release"]
        if self.parallel is not None:
            build_cmd += ["--parallel", str(self.parallel)]
        install_cmd = ["cmake", "--install", str(build_temp), "--prefix", str(install_prefix)]
        if self.cpp_tests:
            # Build every library, extension and C++ test in one CMake
            # invocation so progress is reported against one global total.
            self.spawn(build_cmd)
            self.spawn(install_cmd)
            ctest_cmd = [
                "ctest",
                "--test-dir",
                str(build_temp),
                "--build-config",
                "Release",
                "--output-on-failure",
            ]
            if self.parallel is not None:
                ctest_cmd += ["--parallel", str(self.parallel)]
            self.spawn(ctest_cmd)
        else:
            self.spawn(build_cmd)
            self.spawn(install_cmd)


class BuildBenchmarks(Command):
    """Builds the C++ benchmark executables with RelWithDebInfo and debug symbols."""

    description = "builds C++ benchmark executables (RelWithDebInfo + DWARF debug info)"
    user_options = [
        ("build-temp=", "t", "temporary build directory"),
        ("gprof", None, "add -pg instrumentation for gprof profiling"),
        (
            "with-upstream-onnx",
            None,
            (
                "set -DONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX=ON so CMake fetches "
                "and builds upstream onnx via FetchContent, which makes "
                "BENCH_HAS_UPSTREAM_ONNX defined for bench_load_file"
            ),
        ),
        ("parallel=", "j", "number of parallel build jobs"),
    ]
    boolean_options = ["gprof", "with-upstream-onnx"]

    def initialize_options(self):
        """Initializes default values for command options."""
        self.build_temp = None
        self.gprof = False
        self.with_upstream_onnx = False
        self.parallel = _default_parallel_jobs()

    def finalize_options(self):
        """Sets the build directory to 'build/benchmarks' when not explicitly specified."""
        if self.build_temp is None:
            self.build_temp = os.path.join("build", "benchmarks")

    def run(self):
        """Configures CMake with RelWithDebInfo and ONNX_LIGHT_BUILD_BENCHMARKS=ON.

        Then builds the bench_parse_serialize target.
        """
        root = Path(__file__).resolve().parent
        build_temp = Path(self.build_temp).resolve()
        build_temp.mkdir(parents=True, exist_ok=True)

        cmake_args = [
            "cmake",
            "-S",
            str(root),
            "-B",
            str(build_temp),
            "-DCMAKE_BUILD_TYPE=RelWithDebInfo",
            "-DONNX_LIGHT_BUILD_BENCHMARKS=ON",
            "-DONNX_LIGHT_BUILD_PYTHON=OFF",
        ]
        if self.gprof:
            cmake_args.append("-DONNX_LIGHT_BENCH_GPROF=ON")
        if self.with_upstream_onnx:
            onnx_prefix = _detect_upstream_onnx_prefix()
            if onnx_prefix is not None:
                print(f"--with-upstream-onnx: pip onnx detected at {onnx_prefix}")
            print("--with-upstream-onnx: CMake will FetchContent+build upstream onnx")
            cmake_args.append("-DONNX_LIGHT_BENCH_WITH_UPSTREAM_ONNX=ON")

        self.spawn(cmake_args)
        build_cmd = ["cmake", "--build", str(build_temp), "--target", "bench_parse_serialize"]
        if self.parallel is not None:
            build_cmd += ["--parallel", str(self.parallel)]
        self.spawn(build_cmd)

        bench_bin = build_temp / "bench_parse_serialize"
        if bench_bin.exists():
            print(f"\nBenchmark binary: {bench_bin}")
            print(f"Usage: {bench_bin} -n 20 -t 1\n")


setup(
    name="onnx-light",
    version="0.1.17",
    packages=["onnx_light"],
    distclass=NoConfigDistribution,
    cmdclass={"build_ext": BuildExt, "build_benchmarks": BuildBenchmarks},
)
