"""
QuarterBit Build Script
=======================

Products:
- Axiom:  Memory-efficient quantized optimizer (73% savings)
- Verity: Bit-perfect cross-hardware reproducible optimizer

Compiles Python to binary extensions using Cython for IP protection.
Excludes raw .py source files from wheel - only ships compiled .pyd/.so.

Usage:
    python setup.py build_ext --inplace   # Build locally
    python setup.py bdist_wheel           # Build wheel for PyPI

Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
"""

import os
import sys
import shutil
import platform
from setuptools import setup, find_packages
from setuptools.extension import Extension
from setuptools.command.build_py import build_py

# Check for Cython
try:
    from Cython.Build import cythonize
    from Cython.Distutils import build_ext
    USE_CYTHON = True
except ImportError:
    USE_CYTHON = False
    print("WARNING: Cython not found - source code will NOT be protected!")

# Version
VERSION = "8.0.8"

# Modules to compile and EXCLUDE from wheel (IP protected)
# These .py files will NOT be included - only compiled .pyd/.so binaries ship
PROTECTED_MODULES = [
    "quarterbit/axiom/torch/axiom.py",
    "quarterbit/verity/torch/verity.py",
]

# Custom build_py that excludes protected .py files
class BuildPyExcludeProtected(build_py):
    def find_package_modules(self, package, package_dir):
        modules = super().find_package_modules(package, package_dir)
        # Filter out protected modules - they'll be shipped as .pyd/.so only
        filtered = []
        for (pkg, mod, filepath) in modules:
            full_path = filepath.replace("\\", "/")
            is_protected = any(full_path.endswith(p.replace("quarterbit/", ""))
                              for p in PROTECTED_MODULES)
            if not is_protected:
                filtered.append((pkg, mod, filepath))
            else:
                print(f"  [IP Protected] Excluding source: {filepath}")
        return filtered

# Create extensions if Cython available
extensions = []
if USE_CYTHON:
    for module_path in PROTECTED_MODULES:
        if os.path.exists(module_path):
            module_name = module_path.replace("/", ".").replace("\\", ".").replace(".py", "")
            extensions.append(
                Extension(
                    module_name,
                    [module_path],
                    define_macros=[("CYTHON_TRACE", "0")],
                )
            )

# Cython compiler directives (maximum obfuscation)
cython_directives = {
    "language_level": "3",
    "boundscheck": False,
    "wraparound": False,
    "cdivision": True,
    "embedsignature": False,  # Hide API signatures
}

# Include CUDA libraries and compiled extensions for all products
LIB_FILES = [
    "lib/*.dll",
    "lib/*.so",
    "axiom/cuda/*.dll",
    "axiom/cuda/*.so",
    "axiom/torch/*.pyd",
    "axiom/torch/*.so",
    "verity/cuda/*.dll",
    "verity/cuda/*.so",
    "verity/torch/*.pyd",
    "verity/torch/*.so",
    "q128/*.dll",
    "q128/*.so",
    "q128/*.pyd",
]

# Setup configuration
setup_kwargs = dict(
    name="quarterbit",
    version=VERSION,
    description="Advanced neural network optimizers: Axiom (73% memory savings) + Verity (bit-perfect reproducibility)",
    long_description=open("README.md").read() if os.path.exists("README.md") else "",
    long_description_content_type="text/markdown",
    author="Kyle Clouthier",
    author_email="info@quarterbit.dev",
    url="https://quarterbit.dev",
    packages=find_packages(),
    python_requires=">=3.8",
    install_requires=[
        "torch>=1.8",
        "numpy>=1.20",
    ],
    package_data={
        "quarterbit": LIB_FILES,
    },
    include_package_data=True,
    zip_safe=False,
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "License :: Other/Proprietary License",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: Python :: 3.12",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    keywords="optimizer pytorch deep-learning axiom verity reproducible memory-efficient cross-hardware",
)

# Add Cython compilation and custom build
if USE_CYTHON and extensions:
    setup_kwargs["ext_modules"] = cythonize(
        extensions,
        compiler_directives=cython_directives,
        annotate=False,
    )
    setup_kwargs["cmdclass"] = {
        "build_ext": build_ext,
        "build_py": BuildPyExcludeProtected,
    }
else:
    # Still exclude .py even without Cython (but warn)
    setup_kwargs["cmdclass"] = {"build_py": BuildPyExcludeProtected}

setup(**setup_kwargs)
