#!/usr/bin/env python
"""setup.py script for py_ballisticcalc library"""

import os
import platform
from setuptools import setup, Extension
from pathlib import Path

try:
    from Cython.Build import cythonize
except ImportError:
    import sys

    # use this command to skip building wheel and compiling modules on unsupported platforms
    # pip install --no-build-isolation --no-binary :all: py-ballisticcalc.exts
    setup()
    sys.exit(0)  # Stop installation

_ENV_VAR_IS_ON = {
    "1",
    "true",
    "yes",
    "on",
}


def env_var_is_enabled(var: str):
    return os.environ.get(var, "0").lower() in _ENV_VAR_IS_ON


DISABLE_STRIP = env_var_is_enabled("DISABLE_SHARED_STRIP")

ENABLE_CYTHON_COVERAGE = env_var_is_enabled("CYTHON_COVERAGE")
ENABLE_CYTHON_SAFETY = env_var_is_enabled("CYTHON_SAFETY")

# When coverage is requested, force Cython to regenerate C code to avoid stale builds
CYTHON_FORCE_REGEN = env_var_is_enabled("CYTHON_FORCE_REGEN")

compiler_directives = {
    "language_level": 3,
    "embedsignature": True,
    "cdivision": True,  # use with caution! can affect on "/" operator
    "binding": False,  # Keep as False unless you explicitly want Python-level binding methods
    "c_api_binop_methods": True,
    # Suppress noisy warnings about Python-level names (PreferredUnits, Vector, etc.)
    # that are valid Python imports but flagged by Cython's static analyzer.
    "warn.undeclared": False,
    "warn.unreachable": True,
    "warn.maybe_uninitialized": True,
    # Unused warnings are extremely noisy due to autogenerated pickle/state code; disable them.
    "warn.unused": False,
    "warn.unused_arg": False,
    "warn.multiple_declarators": True,
    "show_performance_hints": True,
    # "emit_code_comments": False,
    "freethreading_compatible": True,
}

# Enable line tracing for Cython coverage when requested
if ENABLE_CYTHON_COVERAGE:
    compiler_directives["linetrace"] = True

# When safety is requested, favor correctness checks over speed
if ENABLE_CYTHON_SAFETY:
    compiler_directives.update(
        {
            "boundscheck": True,
            "wraparound": True,
            "initializedcheck": True,
            "nonecheck": True,
            "cdivision": False,
            "overflowcheck": True,
            "embedsignature": True,
        }
    )

FORCE_CYTHON_MACROS = [("__CYTHON__", "1")]

EXTENSIONS_BASE_DIR = Path("py_ballisticcalc_exts")
SRC_DIR_PATH = EXTENSIONS_BASE_DIR / "src"
INCLUDE_DIR_PATH = EXTENSIONS_BASE_DIR / "include"

# Use absolute paths for include directories
include_dirs = [
    str(EXTENSIONS_BASE_DIR),  # For .pxd files
    str(SRC_DIR_PATH),  # For source-level headers
    str(INCLUDE_DIR_PATH),  # For public headers
]

# Define all C source files and their paths
SOURCE_PATHS = {
    # C++ Sources:
    "interp": SRC_DIR_PATH / "interp.cpp",
    "types": SRC_DIR_PATH / "base_types.cpp",
    "bind": SRC_DIR_PATH / "py_bind.cpp",
    "traj_data": SRC_DIR_PATH / "traj_data.cpp",
    "traj_filter": SRC_DIR_PATH / "traj_filter.cpp",
    "engine": SRC_DIR_PATH / "engine.cpp",
    "euler": SRC_DIR_PATH / "euler.cpp",
    "rk4": SRC_DIR_PATH / "rk4.cpp",
}

# Define dependencies for each extension as a dictionary
# Keys are extension names (as in extension_names list)
# Values are lists of C source file keys from SOURCE_PATHS that they depend on.

_BASE_TYPES_DEPS = set(["types"])
_INTERP_DEPS = set(["interp"])
_TRAJ_DATA_DEPS = set([*_BASE_TYPES_DEPS, *_INTERP_DEPS, "traj_data"])
_BIND_DEPS = set([*_BASE_TYPES_DEPS, "bind"])
_ENGINE_DEPS = set([*_BIND_DEPS, *_TRAJ_DATA_DEPS, "traj_filter", "engine"])
_RK4_DEPS = set([*_ENGINE_DEPS, "rk4"])
_EULER_DEPS = set([*_ENGINE_DEPS, "euler"])
_TEST_DEPS = set([*_ENGINE_DEPS, *_RK4_DEPS, *_EULER_DEPS])

C_EXTENSION_DEPS = {
    # Test modules (expose internal C functions for tests only)
}

CPP_EXTENSION_DEPS = {
    "bind": _BIND_DEPS,
    "traj_data": _TRAJ_DATA_DEPS,
    "base_engine": _ENGINE_DEPS,
    "rk4_engine": _RK4_DEPS,
    "euler_engine": _EULER_DEPS,
    # Test modules (expose internal C++ functions for tests only)
    "_test_helpers": _TEST_DEPS,
    "_test_engine": _TEST_DEPS,
}

TEST_EXTENSIONS_DEPS = {}

# Platform-specific compiler flags
is_msvc = platform.system() == "Windows"
is_macos = platform.system() == "Darwin"


if is_msvc:
    # MSVC-specific flags
    c_compile_args = ["/O2", "/W3"]
    cpp_compile_args = ["/O2", "/W3"]  # "/std:c++11" flag is deprecated
    cpp_extra_link_args = []
    # Crucial for MSVC on ARM
    if platform.machine().startswith("ARM"):
        c_compile_args.append("/fp:precise")
        cpp_compile_args.append("/fp:precise")
elif is_macos:
    c_compile_args = ["-g", "-O0", "-std=c99"]
    cpp_compile_args = ["-O2", "-Wall"]  # assumes it uses -std=c++14 or newer
    cpp_extra_link_args = ["-stdlib=libc++"]
    os.environ["CC"] = "clang"
    os.environ["CXX"] = "clang++"
else:
    # GCC/Clang flags
    c_compile_args = ["-g", "-O0", "-std=c99"]
    cpp_compile_args = ["-fopenmp", "-x", "c++", "-std=c++11", "-O2", "-Wall", "-g"]
    if DISABLE_STRIP:
        cpp_extra_link_args = []
    else:
        # c_compile_args = ["-O3", "-std=c99", "-DNDEBUG"]
        # cpp_compile_args = ["-x", "c++", "-std=c++11", "-O3", "-Wall", "-DNDEBUG"]
        cpp_extra_link_args = ["-Wl,-strip-all"]


# Dynamically create extensions for names in extension_names
def collect_extensions(deps: dict[str, Path], path: Path, *, is_cpp: bool = False):
    extensions_local = []
    for name, deps in deps.items():
        # Use subproject-local .pyx paths (these exist during build)
        sources = [(path / name).with_suffix(".pyx")]

        # Add dependent C source files from the EXTENSION_DEPS dictionary
        # Use .get(name, []) to safely get an empty list if an extension has no explicit C dependencies
        for dep_key in deps:
            if dep_key in SOURCE_PATHS:
                sources.append(SOURCE_PATHS[dep_key])
            else:
                print(f"Warning: C source '{dep_key}' not found in C_SOURCES dictionary for extension '{name}'.")

        define_macros = []
        define_macros.extend(FORCE_CYTHON_MACROS)
        if ENABLE_CYTHON_COVERAGE:
            # Enable tracing in both with-GIL and nogil regions
            define_macros.extend([("CYTHON_TRACE", "1"), ("CYTHON_TRACE_NOGIL", "1")])

        sources = [str(s) for s in sources]

        if not is_cpp:
            extensions_local.append(
                Extension(
                    "py_ballisticcalc_exts." + name,
                    sources=sources,
                    include_dirs=include_dirs,
                    language="c",
                    define_macros=define_macros,
                    extra_compile_args=c_compile_args,
                    # libraries=['m'] # For Linux/macOS, add 'm' for math functions. For Windows, usually not needed or part of default C runtime.
                )
            )
        else:
            extensions_local.append(
                Extension(
                    "py_ballisticcalc_exts." + name,
                    sources=sources,
                    include_dirs=include_dirs,
                    # extra_objects=[],
                    language="c++",
                    define_macros=define_macros,
                    extra_compile_args=cpp_compile_args,
                    extra_link_args=cpp_extra_link_args,
                )
            )

    return extensions_local


# Initialize extensions list
extensions_list = []
extensions_list.extend(collect_extensions(C_EXTENSION_DEPS, EXTENSIONS_BASE_DIR))
extensions_list.extend(collect_extensions(CPP_EXTENSION_DEPS, EXTENSIONS_BASE_DIR, is_cpp=True))

# Standard cythonize with a clean in-project build dir; annotate only if coverage requested
extensions = cythonize(
    extensions_list,
    compiler_directives=compiler_directives,
    annotate=True,  # ENABLE_CYTHON_COVERAGE, # whether to generate .html annotations
    build_dir="build",  # to keep built data
    force=ENABLE_CYTHON_COVERAGE or CYTHON_FORCE_REGEN,
)

setup(ext_modules=extensions)
