"""
Build script for sgnl-cpu-interp with multi-architecture SIMD support.

Builds architecture-specific kernels with appropriate compiler flags:
- x86_64: SSE2, SSE4.1, AVX, AVX2+FMA, AVX-512
- ARM64: NEON
- All: Scalar fallback

Runtime CPU detection selects the best available implementation.
"""

import os
import platform
import sys

import numpy as np
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext

# Detect architecture
machine = platform.machine().lower()
is_x86 = machine in ("x86_64", "amd64", "x86", "i386", "i686")
is_arm = machine in ("arm64", "aarch64", "armv8")
is_windows = sys.platform == "win32"
is_macos = sys.platform == "darwin"

# Base include directories
include_dirs = [
    np.get_include(),
    "csrc",
]

# Base compile args (all platforms)
if is_windows:
    base_compile_args = ["/O2", "/fp:fast"]
    base_link_args: list[str] = []
else:
    base_compile_args = ["-O3", "-ffast-math", "-Wall"]
    base_link_args = []

# Per-file SIMD flags (Unix only - Windows uses /arch:*)
SIMD_FLAGS_UNIX = {
    "convolve_sse2.c": ["-msse2"],
    "convolve_sse4.c": ["-msse4.1"],
    "convolve_avx.c": ["-mavx"],
    "convolve_avx2.c": ["-mavx2", "-mfma"],
    "convolve_avx512.c": ["-mavx512f", "-mfma"],
}

SIMD_FLAGS_MSVC = {
    "convolve_avx.c": ["/arch:AVX"],
    "convolve_avx2.c": ["/arch:AVX2"],
    "convolve_avx512.c": ["/arch:AVX512"],
    # SSE2 is default on x64 MSVC
}

# Source files
common_sources = [
    "csrc/cpu_detect.c",
    "csrc/dispatch.c",
    "csrc/resample_ext_simd.c",
    "csrc/kernels/convolve_scalar.c",
]

# Architecture-specific sources
if is_x86:
    arch_sources = [
        "csrc/kernels/convolve_sse2.c",
        "csrc/kernels/convolve_sse4.c",
        "csrc/kernels/convolve_avx.c",
        "csrc/kernels/convolve_avx2.c",
        "csrc/kernels/convolve_avx512.c",
    ]
elif is_arm:
    arch_sources = [
        "csrc/kernels/convolve_neon.c",
    ]
else:
    arch_sources = []

all_sources = common_sources + arch_sources


class BuildExtSIMD(build_ext):
    """
    Custom build_ext that applies per-file compiler flags for SIMD.

    This allows compiling different source files with different -m flags
    (e.g., -mavx2 for AVX2 files, -msse2 for SSE2 files).
    """

    def build_extension(self, ext):
        if ext.name != "sgnl_cpu_interp._core":
            # Use default behavior for other extensions
            return super().build_extension(ext)

        # Get the appropriate SIMD flags dict
        if is_windows:
            simd_flags = SIMD_FLAGS_MSVC
        else:
            simd_flags = SIMD_FLAGS_UNIX

        # Store original _compile method
        original_compile = self.compiler._compile

        def custom_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
            """Custom compile that adds per-file SIMD flags."""
            # Get basename of source file
            basename = os.path.basename(src)

            # Add SIMD flags if this file needs them
            if basename in simd_flags:
                extra_postargs = list(extra_postargs) + simd_flags[basename]

            return original_compile(obj, src, ext, cc_args, extra_postargs, pp_opts)

        # Monkey-patch the compile method
        self.compiler._compile = custom_compile

        try:
            super().build_extension(ext)
        finally:
            # Restore original compile method
            self.compiler._compile = original_compile


# Libraries (no FFTW or GSL dependency!)
if is_windows:
    libraries = []
else:
    libraries = ["m"]

# Define the extension
ext_simd = Extension(
    "sgnl_cpu_interp._core",
    sources=all_sources,
    include_dirs=include_dirs,
    libraries=libraries,
    extra_compile_args=base_compile_args,
    extra_link_args=base_link_args,
)

setup(
    ext_modules=[ext_simd],
    cmdclass={"build_ext": BuildExtSIMD},
)
