#!/usr/bin/env python
import os
import sys
import shutil
from setuptools import setup, Extension
from Cython.Build import cythonize


def build_extensions(target_dirs):
    """
    Compiles all .py files in the target directories (recursively) into .so/.pyd extensions.
    Args:
        target_dirs (list or str): List of directories or single directory path to scan.
    """
    # Ensure input is a list
    if isinstance(target_dirs, str):
        target_dirs = [target_dirs]

    extensions = []
    project_root = os.getcwd()

    print(f"Scanning directories: {target_dirs}")

    for target_dir in target_dirs:
        # Ensure target directory exists
        if not os.path.exists(target_dir):
            print(f"Warning: Path '{target_dir}' not found. Skipping.")
            continue

        # Get the absolute path of the target
        abs_target_path = os.path.abspath(target_dir)

        # Check if the target is a file or directory
        if os.path.isfile(abs_target_path):
            if abs_target_path.endswith(".py"):
                file_path = abs_target_path
                rel_path = os.path.relpath(file_path, project_root)
                module_name = os.path.splitext(rel_path)[0].replace(os.sep, ".")
                print(f"Found file: {rel_path} -> {module_name}")
                extensions.append(Extension(module_name, [file_path]))
            continue

        print(f"Scanning {abs_target_path} for Python files...")

        # Walk through the directory
        for root, dirs, files in os.walk(abs_target_path):
            for file in files:
                if file.endswith(".py"):
                    file_path = os.path.join(root, file)

                    # Skip this script if it happens to be in the target dir
                    if os.path.abspath(file_path) == os.path.abspath(__file__):
                        continue

                    # Skip setup.py if it exists
                    if file == "setup.py":
                        continue

                    # Determine the module name based on path relative to project root
                    # This ensures imports like 'from app.services import ...' work
                    try:
                        rel_path = os.path.relpath(file_path, project_root)
                    except ValueError:
                        # If file is not under project root, we can't easily determine module name
                        # relative to project root. Skip or warn.
                        print(
                            f"Skipping {file_path}: cannot determine relative path to {project_root}"
                        )
                        continue

                    # Convert file path to module name (e.g. app/services/foo.py -> app.services.foo)
                    module_name = os.path.splitext(rel_path)[0].replace(os.sep, ".")

                    print(f"Found: {rel_path} -> {module_name}")

                    extensions.append(Extension(module_name, [file_path]))

    if not extensions:
        print("No Python files found to compile.")
        return

    print(f"\nCompiling {len(extensions)} modules...")

    # Build options
    # compiler_directives: language_level=3 for Python 3
    # force=True: force recompilation even if timestamps are up to date
    try:
        setup(
            ext_modules=cythonize(
                extensions,
                compiler_directives={"language_level": "3"},
                build_dir="build",  # Put intermediate files in build/ directory
                force=True,
            ),
            script_args=["build_ext", "--inplace"],
        )
        print("\nCompilation successful!")
    except Exception as e:
        print(f"\nCompilation failed: {e}")
        sys.exit(1)
    finally:
        # Cleanup build directory (intermediate C files)
        if os.path.exists("build"):
            print("Cleaning up build artifacts...")
            try:
                shutil.rmtree("build")
            except OSError as e:
                print(f"Warning: Failed to clean up build directory: {e}")


if __name__ == "__main__":
    # Default directories to compile if none provided
    DEFAULT_TARGETS = [
        "app/services",
        "app/native/wndb",
        "app/algorithms",
        "app/infra/epanet/epanet.py",
    ]

    # Check for help flag
    if len(sys.argv) > 1 and sys.argv[1] in ["--help", "-h"]:
        print("Usage: python scripts/build_extensions.py [directory1] [directory2] ...")
        print(
            "Compiles all Python files in the given directories into C extensions (.so/.pyd)."
        )
        print(f"Default directories if none provided: {DEFAULT_TARGETS}")
        sys.exit(0)

    if len(sys.argv) > 1:
        target_directories = sys.argv[1:]
    else:
        print(f"No directories specified. Using defaults: {DEFAULT_TARGETS}")
        target_directories = DEFAULT_TARGETS

    build_extensions(target_directories)