import os
import shutil
import subprocess
import zipfile  # <-- Add this import at the top
from pathlib import Path

def build_and_clean():
    print("🧠 Generating IDE Stub Files (.pyi)...")
    try:
        subprocess.run(["stubgen", "--parse-only", "./mxpic", "-o", "."], check=True)
        print("✅ Stubs generated successfully.")
    except subprocess.CalledProcessError:
        print("❌ Failed to generate stubs.")
        return

    print("🚀 Starting Wheel Build Process...")
    try:
        subprocess.run(["python", "-m", "build", "--wheel"], check=True)
        print("✅ Base Wheel built successfully.")
    except subprocess.CalledProcessError:
        print("❌ Build failed. Stopping cleanup.")
        return

    # --- THE SURGICAL INJECTION ---
    print("💉 Injecting .pyi stubs directly into the Wheel...")
    wheel_files = list(Path("dist").glob("*.whl"))
    
    if wheel_files:
        target_wheel = wheel_files[0]
        # Open the wheel file in 'append' mode
        with zipfile.ZipFile(target_wheel, 'a', compression=zipfile.ZIP_DEFLATED) as wheel_zip:
            injected_count = 0
            for pyi_file in Path("mxpic").rglob("*.pyi"):
                # Inject the file, preserving the 'mxpic/...' folder structure
                wheel_zip.write(pyi_file, arcname=pyi_file)
                injected_count += 1
        print(f"✅ Successfully injected {injected_count} stub files into {target_wheel.name}")
    else:
        print("❌ Could not find the .whl file in the dist/ directory.")

    # --- THE CLEANUP ---
    print("\n🧹 Commencing Workspace Cleanup...")
    
    src_path = Path("mxpic")
    if src_path.exists():
        # Now it is safe to delete the stubs from the source folder!
        for pyi_file in src_path.rglob("*.pyi"):
            pyi_file.unlink()

        # Clean Cython artifacts
        for c_file in src_path.rglob("*.c"):
            c_file.unlink()

    # Delete the bulky build directories
    folders_to_remove = ["build", "mxpic.egg-info"]
    for folder in folders_to_remove:
        folder_path = Path(folder)
        if folder_path.exists() and folder_path.is_dir():
            shutil.rmtree(folder_path)

    print("\n✨ Workspace is clean. Your ultimate .whl is ready in dist/!")

if __name__ == "__main__":
    build_and_clean()