import os
import glob
import shutil
import subprocess
import sys
import argparse
import compileall

# Step 1: Compile .py files to .pyc files
def compile_files_to_pyc(root_dir):
    compileall.compile_dir(root_dir, force=True, legacy=True)
    for py_file in glob.glob(f"{root_dir}/**/*.py", recursive=True):
        py_cache_dir = os.path.join(os.path.dirname(py_file), '__pycache__')
        for pyc_file in glob.glob(f"{py_cache_dir}/*.pyc"):
            target_file = os.path.join(os.path.dirname(py_file), os.path.basename(pyc_file))
            shutil.move(pyc_file, target_file)
            print(f"Moved {pyc_file} to {target_file}")
        if os.path.exists(py_cache_dir):
            shutil.rmtree(py_cache_dir)

# Step 2: Remove .py files from the build directory except setup.py
def remove_py_files_from_build(build_dir):
    for py_file in glob.glob(f"{build_dir}/**/*.py", recursive=True):
        if os.path.basename(py_file) == "setup.py":
            continue
        os.remove(py_file)
        print(f"Removed {py_file} from build directory")

# Step 3: Backup .py files
def backup_py_files(root_dir, backup_dir):
    if os.path.exists(backup_dir):
        shutil.rmtree(backup_dir)
    shutil.copytree(root_dir, backup_dir, ignore=shutil.ignore_patterns('__pycache__', 'build', 'dist', '*.egg-info', '.git'))

# Step 4: Restore .py files from backup
def restore_py_files_from_backup(root_dir, backup_dir):
    for item in os.listdir(backup_dir):
        s = os.path.join(backup_dir, item)
        d = os.path.join(root_dir, item)
        if os.path.isdir(s):
            shutil.copytree(s, d, dirs_exist_ok=True)
        else:
            shutil.copy2(s, d)

# Step 5: Setup and compile
def compile_extensions(build_dir):
    subprocess.run([sys.executable, "setup.py", "bdist_wheel", "--dist-dir", "../dist"], cwd=build_dir)

def main(include_source):
    root_dir = "."
    backup_dir = "./backup"
    build_dir = "./build"

    if not include_source:
        backup_py_files(root_dir, backup_dir)
        compile_files_to_pyc(root_dir)

        # Copy files to build_dir for packaging
        if os.path.exists(build_dir):
            shutil.rmtree(build_dir)
        shutil.copytree(root_dir, build_dir, ignore=shutil.ignore_patterns('__pycache__', 'build', 'dist', '*.egg-info', '.git'))

        remove_py_files_from_build(build_dir)

    compile_extensions(build_dir)

    if not include_source:
        restore_py_files_from_backup(root_dir, backup_dir)
        shutil.rmtree(backup_dir)
        if os.path.exists(build_dir):
            shutil.rmtree(build_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build and package the Python project.")
    parser.add_argument("--include-source", action="store_true", help="Include source code in the package")
    args = parser.parse_args()

    main(args.include_source)