"""
Setup script for ciffy C extension.

Metadata is defined in pyproject.toml. This file only handles:
1. C extension compilation
2. Hash table generation before build (downloads CCD if needed)
"""

from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
from setuptools.command.sdist import sdist
import os
import sys
import subprocess
import shutil
import gzip
import numpy


# URL for the PDB Chemical Component Dictionary
CCD_URL = "https://files.wwpdb.org/pub/pdb/data/monomers/components.cif.gz"


def download_ccd(dest_path):
    """Download and decompress the CCD file."""
    import urllib.request

    print(f"Downloading CCD from {CCD_URL}...")
    gz_path = dest_path + ".gz"

    try:
        urllib.request.urlretrieve(CCD_URL, gz_path)
        print("Decompressing CCD...")
        with gzip.open(gz_path, 'rb') as f_in:
            with open(dest_path, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(gz_path)
        print(f"CCD downloaded to {dest_path}")
        return True
    except Exception as e:
        print(f"Failed to download CCD: {e}")
        if os.path.exists(gz_path):
            os.remove(gz_path)
        return False


def get_ccd_path():
    """Get path to CCD file, downloading if necessary."""
    # Check environment variable first
    ccd_path = os.environ.get("CIFFY_CCD_PATH")
    if ccd_path and os.path.exists(ccd_path):
        return ccd_path

    # Use centralized cache location
    cache_dir = os.path.expanduser("~/.cache/ciffy")
    ccd_path = os.path.join(cache_dir, "components.cif")

    if os.path.exists(ccd_path):
        return ccd_path

    # Download to cache directory
    os.makedirs(cache_dir, exist_ok=True)
    if download_ccd(ccd_path):
        return ccd_path

    return None


def generate_hash_tables(force=False):
    """Run the hash table generator.

    Args:
        force: If True, regenerate even if files exist (for sdist builds)
    """
    hash_dir = os.path.join(os.path.dirname(__file__), 'ciffy', 'src', 'hash')
    biochem_dir = os.path.join(os.path.dirname(__file__), 'ciffy', 'biochemistry')

    # All files generated by codegen/generate.py
    required_hash_files = [
        'atom.c', 'element.c', 'entity.c', 'ion.c', 'molecule.c', 'residue.c', 'reverse.h'
    ]
    required_python_files = [
        '_generated_atoms.py', '_generated_dihedraltypes.py',
        '_generated_elements.py', '_generated_molecule.py', '_generated_residues.py'
    ]

    def check_generated_files_exist():
        """Check if all required generated files exist."""
        hash_ok = all(os.path.exists(os.path.join(hash_dir, f)) for f in required_hash_files)
        python_ok = all(os.path.exists(os.path.join(biochem_dir, f)) for f in required_python_files)
        return hash_ok and python_ok

    def get_missing_files():
        """Get list of missing generated files."""
        missing = []
        for f in required_hash_files:
            if not os.path.exists(os.path.join(hash_dir, f)):
                missing.append(f'src/hash/{f}')
        for f in required_python_files:
            if not os.path.exists(os.path.join(biochem_dir, f)):
                missing.append(f'biochemistry/{f}')
        return missing

    codegen_dir = os.path.join(os.path.dirname(__file__), 'codegen')
    if not os.path.exists(codegen_dir):
        print("Warning: codegen package not found, skipping hash generation")
        return

    # Check if generated files already exist (users installing from PyPI)
    if check_generated_files_exist() and not force:
        print("Using pre-generated files")
        return

    # Need to generate - get CCD file
    ccd_path = get_ccd_path()
    if not ccd_path:
        if check_generated_files_exist():
            print("Warning: CCD not available, using existing generated files")
            return
        else:
            print("ERROR: CCD file required but not found. Set CIFFY_CCD_PATH or download from:")
            print(f"  {CCD_URL}")
            return

    # Check if gperf is available (need 3.1+ for constants-prefix)
    gperf_path = None
    for path in ["/opt/homebrew/bin/gperf", "/usr/local/bin/gperf"]:
        if os.path.exists(path):
            gperf_path = path
            break
    if gperf_path is None:
        gperf_path = shutil.which("gperf")

    args = [ccd_path]
    if gperf_path is None:
        # gperf not found - check if pre-generated files exist
        if check_generated_files_exist():
            print("Warning: gperf not found, using pre-generated files")
            return
        else:
            # Fatal error - can't build without gperf or pre-generated files
            missing = get_missing_files()
            print("\n" + "=" * 70)
            print("ERROR: Cannot build ciffy - gperf is required but not installed")
            print("=" * 70)
            print("\nThe C extension requires generated files that are created using gperf.")
            print("Pre-generated files are not available.")
            print(f"\nMissing files ({len(missing)}):")
            for f in missing[:10]:  # Show first 10
                print(f"  - {f}")
            if len(missing) > 10:
                print(f"  ... and {len(missing) - 10} more")
            print("\nTo fix this, install gperf:")
            print("  - Ubuntu/Debian: sudo apt install gperf")
            print("  - Fedora/RHEL:   sudo dnf install gperf")
            print("  - macOS:         brew install gperf")
            print("  - Conda:         conda install -c conda-forge gperf")
            print("\nThen re-run: pip install -e .")
            print("=" * 70 + "\n")
            sys.exit(1)
    else:
        args.extend(["--gperf-path", gperf_path])

    print("Generating hash lookup tables...")
    env = os.environ.copy()
    # Build PYTHONPATH that includes:
    # 1. Project directory (for codegen module)
    # 2. Current sys.path (includes pip's isolated build env with dependencies)
    # 3. Any existing PYTHONPATH
    project_dir = os.path.dirname(__file__)
    paths = [project_dir] + sys.path
    existing_pythonpath = env.get("PYTHONPATH", "")
    if existing_pythonpath:
        paths.append(existing_pythonpath)
    env["PYTHONPATH"] = os.pathsep.join(p for p in paths if p)

    # Run as module to support relative imports
    result = subprocess.run(
        [sys.executable, "-m", "codegen.generate"] + args,
        env=env,
        cwd=os.path.dirname(__file__),
        capture_output=True,
        text=True
    )

    if result.returncode != 0:
        print(f"Warning: Hash generation failed: {result.stderr}")
    else:
        print(result.stdout)

    # Final check - ensure all generated files were created
    if not check_generated_files_exist():
        missing = get_missing_files()
        stderr = result.stderr if result.returncode != 0 else ""
        print("\n" + "=" * 70)
        print("ERROR: Code generation failed")
        print("=" * 70)
        print(f"\nMissing files ({len(missing)}):")
        for f in missing:
            print(f"  - {f}")

        # Provide targeted error message based on failure mode
        if "ModuleNotFoundError" in stderr or "No module named" in stderr:
            # Extract module name from error
            import re
            match = re.search(r"No module named ['\"]?(\w+)['\"]?", stderr)
            module_name = match.group(1) if match else "unknown"
            print(f"\nA required Python module is missing: {module_name}")
            print("This is a build dependency issue. Install the missing module:")
            print(f"  pip install {module_name}")
            print("\nIf you're installing ciffy, this indicates a packaging bug.")
            print("Please report this issue at: https://github.com/hmblair/ciffy/issues")
        elif "gperf" in stderr.lower() or not gperf_path:
            print("\nThis usually means gperf failed. Check that gperf is installed correctly:")
            print("  gperf --version")
            print("\nTo install gperf:")
            print("  - Ubuntu/Debian: sudo apt install gperf")
            print("  - Fedora/RHEL:   sudo dnf install gperf")
            print("  - macOS:         brew install gperf")
            print("  - Conda:         conda install -c conda-forge gperf")
        else:
            print("\nCode generation failed. See error above for details.")
            if stderr:
                print(f"\nError output:\n{stderr[:500]}")

        print("=" * 70 + "\n")
        sys.exit(1)


class GenerateAndBuildPy(build_py):
    """Custom build_py that generates Python files before copying.

    This runs BEFORE build_ext, ensuring generated Python files
    (_generated_*.py) are included in wheel distributions.
    """

    def run(self):
        generate_hash_tables(force=False)
        super().run()


class GenerateAndBuildExt(build_ext):
    """Custom build_ext that generates hash tables before compiling."""

    def run(self):
        # Codegen may have already run in build_py, but check again for safety
        generate_hash_tables(force=False)
        super().run()

        # On macOS, fix up libomp linkage to use @rpath instead of absolute path
        # This avoids conflicts when PyTorch (with its own libomp) is also loaded
        if sys.platform == 'darwin':
            self._fix_libomp_rpath()

    def _fix_libomp_rpath(self):
        """Change libomp reference from absolute path to @rpath on macOS."""
        for ext in self.extensions:
            ext_path = self.get_ext_fullpath(ext.name)
            if not os.path.exists(ext_path):
                continue

            # Check if this extension links against libomp
            try:
                result = subprocess.run(
                    ['otool', '-L', ext_path],
                    capture_output=True, text=True, timeout=10
                )
                if '/libomp.dylib' not in result.stdout:
                    continue

                # Find the absolute path to libomp
                for line in result.stdout.split('\n'):
                    if '/libomp.dylib' in line and '@rpath' not in line:
                        # Extract the path (first part before ' (')
                        old_path = line.strip().split(' (')[0]
                        if old_path:
                            print(f"Fixing libomp linkage: {old_path} -> @rpath/libomp.dylib")
                            subprocess.run(
                                ['install_name_tool', '-change', old_path,
                                 '@rpath/libomp.dylib', ext_path],
                                capture_output=True, timeout=10
                            )
                            break
            except Exception as e:
                print(f"Warning: Could not fix libomp rpath: {e}")


class GenerateAndSdist(sdist):
    """Custom sdist that ensures hash tables are generated before packaging."""

    def run(self):
        # Force regeneration for sdist to ensure latest definitions
        generate_hash_tables(force=True)
        super().run()


# ============================================================================
# Compiler configuration
# ============================================================================

def check_openmp_available():
    """
    Check if OpenMP is available by trying to compile a simple test program.

    Returns:
        Tuple of (compile_args, link_args) or (None, None) if not available
    """
    import tempfile

    if sys.platform == 'darwin':
        # macOS: need libomp from Homebrew, conda, or torch
        compile_args = ['-Xpreprocessor', '-fopenmp']
        link_args = []
        libomp_path = None

        # CRITICAL FIX FOR DUAL-LIBOMP CONFLICT:
        # On macOS, environments often have multiple libomp copies (conda's system + torch's bundled).
        # When both try to initialize, we get "OMP: Error #15: already initialized".
        # Solution: Prioritize torch's libomp since torch will load it anyway.
        # This ensures we use the SAME libomp that torch will load at runtime.

        # Priority 1: Use torch's libomp (avoids conflict with conda's system libomp)
        # CRITICAL: Use absolute path, not just -L flag, to ensure torch's libomp
        # is loaded instead of conda's system libomp in the same environment
        try:
            import torch
            torch_lib = os.path.join(os.path.dirname(torch.__file__), 'lib')
            torch_libomp = os.path.join(torch_lib, 'libomp.dylib')
            if os.path.exists(torch_libomp):
                # Use absolute path to torch's libomp to ensure it takes precedence
                # over conda's system libomp in the same environment
                link_args.append(torch_libomp)  # Link directly to torch's libomp
                libomp_path = torch_lib
                print(f"Using torch's libomp (absolute path): {torch_libomp}")
        except ImportError:
            pass

        # Priority 2: Check target conda environment's libomp (if not using torch)
        if not libomp_path:
            target_env_prefix = os.environ.get('CONDA_PREFIX')
            if target_env_prefix and target_env_prefix != sys.prefix:
                target_libomp = os.path.join(target_env_prefix, 'lib', 'libomp.dylib')
                if os.path.exists(target_libomp):
                    libomp_path = os.path.join(target_env_prefix, 'lib')
                    print(f"Using target conda environment's libomp: {libomp_path}")

        # Priority 3: Use current sys.prefix environment's libomp
        if not libomp_path:
            conda_libomp = os.path.join(sys.prefix, 'lib', 'libomp.dylib')
            if os.path.exists(conda_libomp):
                libomp_path = os.path.join(sys.prefix, 'lib')
                print("Using current environment's libomp")

        # Priority 4: Fall back to Homebrew libomp
        if not libomp_path:
            for homebrew_path in ['/opt/homebrew/opt/libomp/lib', '/usr/local/opt/libomp/lib']:
                if os.path.exists(homebrew_path):
                    libomp_path = homebrew_path
                    print("Using Homebrew libomp")
                    break

        if not libomp_path:
            # libomp not found
            return None, None

        # Set up linking with explicit path to avoid conflicts
        # If we already added absolute path (torch case), don't add again
        if not any(arg.endswith('libomp.dylib') for arg in link_args):
            link_args.append(f'-L{libomp_path}')
            link_args.append('-lomp')

        include_path = libomp_path.replace('/lib', '/include')
        if os.path.exists(include_path):
            compile_args.append(f'-I{include_path}')
    else:
        # Linux/Windows
        compile_args = ['-fopenmp']
        link_args = ['-fopenmp']

    # Try to compile a test program
    test_code = '#include <omp.h>\nint main() { return omp_get_num_threads(); }'
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.c', delete=False) as f:
            f.write(test_code)
            test_file = f.name

        import sysconfig
        cc = os.environ.get('CC', sysconfig.get_config_var('CC') or 'cc')
        # Handle cases where CC might be "cc -pthread" etc.
        cc = cc.split()[0]

        result = subprocess.run(
            [cc] + compile_args + [test_file, '-o', '/dev/null'] + link_args,
            capture_output=True,
            timeout=30
        )
        os.unlink(test_file)
        if result.returncode == 0:
            return compile_args, link_args
    except Exception:
        pass

    return None, None


# Build compile args
extra_compile_args = ['-O3']
extra_link_args = ['-lz']  # zlib for transparent gzip support

# Enable OpenMP unless CIFFY_NO_OPENMP is set
if os.environ.get('CIFFY_NO_OPENMP', '').lower() not in ('1', 'true', 'yes'):
    omp_compile, omp_link = check_openmp_available()
    if omp_compile and omp_link:
        extra_compile_args.extend(omp_compile)
        extra_link_args.extend(omp_link)

        # On macOS, configure rpaths to use the same libomp as the environment
        # This avoids conflicts with numpy, torch, and other packages
        if sys.platform == 'darwin':
            # Add header padding for install_name_tool modifications
            extra_link_args.append('-Wl,-headerpad_max_install_names')

            # Determine which libomp we're using and add matching rpaths
            # This must match the priority order from check_openmp_available()
            rpath_added = False

            # Priority 1: torch's libomp (avoids dual-libomp conflict)
            try:
                import torch
                torch_lib = os.path.join(os.path.dirname(torch.__file__), 'lib')
                torch_libomp = os.path.join(torch_lib, 'libomp.dylib')
                if os.path.exists(torch_libomp):
                    extra_link_args.append(f'-Wl,-rpath,{torch_lib}')
                    print(f"  Added rpath for torch's libomp: {torch_lib}")
                    rpath_added = True
            except ImportError:
                pass

            # Priority 2: Target conda environment's libomp
            if not rpath_added:
                target_env_prefix = os.environ.get('CONDA_PREFIX')
                if target_env_prefix and target_env_prefix != sys.prefix:
                    target_conda_lib = os.path.join(target_env_prefix, 'lib')
                    if os.path.exists(os.path.join(target_conda_lib, 'libomp.dylib')):
                        extra_link_args.append(f'-Wl,-rpath,{target_conda_lib}')
                        print(f"  Added rpath for target conda environment's libomp: {target_conda_lib}")
                        rpath_added = True

            # Priority 3: Current environment's libomp
            if not rpath_added:
                conda_lib = os.path.join(sys.prefix, 'lib')
                if os.path.exists(os.path.join(conda_lib, 'libomp.dylib')):
                    extra_link_args.append(f'-Wl,-rpath,{conda_lib}')
                    print(f"  Added rpath for current environment's libomp: {conda_lib}")
                    rpath_added = True

            # Priority 4: Homebrew libomp
            if not rpath_added:
                for homebrew_path in ['/opt/homebrew/opt/libomp/lib', '/usr/local/opt/libomp/lib']:
                    if os.path.exists(homebrew_path):
                        extra_link_args.append(f'-Wl,-rpath,{homebrew_path}')
                        print(f"  Added rpath for Homebrew libomp: {homebrew_path}")
                        break

        print("OpenMP enabled")
    else:
        print("OpenMP not available (install libomp on macOS: brew install libomp)")
else:
    print("OpenMP disabled via CIFFY_NO_OPENMP")

# ============================================================================
# C extension module
# ============================================================================

c_sources = [
    'ciffy/src/module.c',
    'ciffy/src/pyutils.c',
    # CIF I/O module
    'ciffy/src/cif/io.c',
    'ciffy/src/cif/parser.c',
    'ciffy/src/cif/writer.c',
    'ciffy/src/cif/registry.c',
    'ciffy/src/cif/connections.c',
    # Graph operations
    'ciffy/src/graph/graph.c',
    'ciffy/src/graph/bindings.c',
]

# Validate source files exist
missing_sources = [src for src in c_sources if not os.path.exists(src)]
if missing_sources:
    print(f"ERROR: Missing source files: {missing_sources}")
    print("Make sure you're building from the ciffy root directory.")
    sys.exit(1)

ext_module = Extension(
    name="ciffy._c",
    sources=c_sources,
    include_dirs=[numpy.get_include(), 'ciffy/src'],
    extra_compile_args=extra_compile_args,
    extra_link_args=extra_link_args,
    language='c',  # Explicitly specify C (not C++)
)

# ============================================================================
# Build extensions list
# ============================================================================

ext_modules = [ext_module]
cmdclass = {
    'build_py': GenerateAndBuildPy,
    'build_ext': GenerateAndBuildExt,
    'sdist': GenerateAndSdist,
}

setup(
    ext_modules=ext_modules,
    cmdclass=cmdclass,
)
