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


def delete_source_files(target_dirs):
    """
    Deletes the original .py files in the target directories after successful compilation.
    Caution: This will remove your source code!
    """
    if isinstance(target_dirs, str):
        target_dirs = [target_dirs]

    project_root = os.getcwd()
    print(f"Deleting source .py files in: {target_dirs}")

    for target_dir in target_dirs:
        if not os.path.exists(target_dir):
            continue

        abs_target_path = os.path.abspath(target_dir)

        if os.path.isfile(abs_target_path):
            if abs_target_path.endswith(".py") and os.path.abspath(
                abs_target_path
            ) != os.path.abspath(__file__):
                print(f"Deleting source: {abs_target_path}")
                os.remove(abs_target_path)
            continue

        # If it's a directory, walk through it
        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 itself
                    if os.path.abspath(file_path) == os.path.abspath(__file__):
                        continue
                    # Skip __init__.py if you want to keep package structure,
                    # but usually Cython can handle them too.
                    print(f"Deleting source: {file_path}")
                    os.remove(file_path)


def clean_extensions(target_dirs):
    """
    Deletes all compiled .so/.pyd files in the target directories.
    """
    if isinstance(target_dirs, str):
        target_dirs = [target_dirs]

    project_root = os.getcwd()
    print(f"Cleaning compiled files in: {target_dirs}")

    for target_dir in target_dirs:
        if not os.path.exists(target_dir):
            continue

        abs_target_path = os.path.abspath(target_dir)

        if os.path.isfile(abs_target_path):
            # If it's a .py file, look for its .so/.pyd counterpart
            dir_name = os.path.dirname(abs_target_path)
            base_name = os.path.splitext(os.path.basename(abs_target_path))[0]
            for file in os.listdir(dir_name):
                if file.startswith(base_name) and (
                    file.endswith(".so") or file.endswith(".pyd")
                ):
                    file_path = os.path.join(dir_name, file)
                    print(f"Removing: {file_path}")
                    os.remove(file_path)
            continue

        # If it's a directory, walk through it
        for root, dirs, files in os.walk(abs_target_path):
            for file in files:
                if file.endswith(".so") or file.endswith(".pyd"):
                    file_path = os.path.join(root, file)
                    print(f"Removing: {file_path}")
                    os.remove(file_path)


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 [--clean|--delete-source] [directory1] [directory2] ..."
        )
        print(
            "Compiles all Python files in the given directories into C extensions (.so/.pyd)."
        )
        print("Use --clean to remove existing compiled extensions.")
        print("Use --delete-source to remove source .py files after compilation.")
        print(f"Default directories if none provided: {DEFAULT_TARGETS}")
        sys.exit(0)

    clean_mode = False
    delete_source_mode = False
    if len(sys.argv) > 1:
        if sys.argv[1] == "--clean":
            clean_mode = True
            target_directories = sys.argv[2:]
        elif sys.argv[1] == "--delete-source":
            delete_source_mode = True
            target_directories = sys.argv[2:]
        else:
            target_directories = sys.argv[1:]
    else:
        print(f"No directories specified. Using defaults: {DEFAULT_TARGETS}")
        target_directories = DEFAULT_TARGETS

    if clean_mode:
        clean_extensions(target_directories)
    elif delete_source_mode:
        delete_source_files(target_directories)
    else:
        build_extensions(target_directories)
