build.py

"""
Build script for creating Transcribair executable.
Handles PyInstaller build with proper configuration.

Usage:
    python build.py                  # Build with CUDA DLLs
    python build.py --no-cuda        # Build without CUDA (CPU-only, smaller)
    python build.py --zip            # Build and create distribution ZIP
    python build.py --no-cuda --zip  # CPU-only ZIP distribution
"""
import argparse
import subprocess
import sys
import time
import zipfile
from pathlib import Path
import shutil
import os


def check_requirements():
    """Check if required tools are installed."""
    try:
        import PyInstaller
        print("[OK] PyInstaller found")
    except ImportError:
        print("[ERROR] PyInstaller not found")
        print("  Install with: pip install pyinstaller")
        return False

    return True


def get_size_mb(path: Path) -> float:
    """Get file or directory size in MB."""
    if path.is_file():
        return path.stat().st_size / (1024 * 1024)
    elif path.is_dir():
        total = sum(f.stat().st_size for f in path.rglob('*') if f.is_file())
        return total / (1024 * 1024)
    return 0.0


def copy_cuda_dlls(output_dir: Path):
    """Copy CUDA/cuDNN DLLs from nvidia pip packages into the build output.

    Searches the current Python environment's site-packages for nvidia
    packages (cudnn, cublas, cuda_nvrtc) and copies their DLLs alongside
    the executable so Windows DLL search finds them naturally.
    """
    print("\nCopying CUDA/cuDNN DLLs...")

    # Find nvidia packages in site-packages
    nvidia_dir = None
    for p in sys.path:
        candidate = Path(p) / "nvidia"
        if candidate.is_dir():
            nvidia_dir = candidate
            break

    if nvidia_dir is None:
        print("[WARNING] No nvidia packages found in site-packages.")
        print("         CUDA DLLs will not be bundled. GPU acceleration")
        print("         will only work if CUDA is installed system-wide.")
        return 0

    # Libraries to copy DLLs from
    nvidia_libs = ["cudnn", "cublas", "cuda_nvrtc"]
    copied = 0

    for lib_name in nvidia_libs:
        bin_dir = nvidia_dir / lib_name / "bin"
        if not bin_dir.is_dir():
            # Try alternate directory structures (e.g. nvidia_cublas_cu12)
            for d in nvidia_dir.iterdir():
                if d.is_dir() and lib_name in d.name:
                    alt_bin = d / "bin"
                    if alt_bin.is_dir():
                        bin_dir = alt_bin
                        break

        if not bin_dir.is_dir():
            print(f"  [WARNING] nvidia/{lib_name}/bin/ not found, skipping")
            continue

        for dll in bin_dir.glob("*.dll"):
            dest = output_dir / dll.name
            if not dest.exists():
                shutil.copy2(dll, dest)
                size_mb = dll.stat().st_size / (1024 * 1024)
                print(f"  [OK] {dll.name} ({size_mb:.1f} MB)")
                copied += 1

    if copied > 0:
        print(f"\n[OK] Copied {copied} CUDA DLLs into build output")
    else:
        print("[WARNING] No CUDA DLLs found to copy")

    return copied


def create_zip(output_dir: Path, include_cuda: bool):
    """Create a distribution ZIP from the build output."""
    from __version__ import __version__

    variant = "CUDA" if include_cuda else "CPU"
    zip_name = f"TranscribAIr-{__version__}-{variant}.zip"
    zip_path = output_dir.parent / zip_name

    print(f"\nCreating distribution: {zip_name}")

    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        for file in output_dir.rglob('*'):
            if file.is_file():
                arcname = Path("TranscribAIr") / file.relative_to(output_dir)
                zf.write(file, arcname)

    zip_size = get_size_mb(zip_path)
    print(f"[OK] Created {zip_path} ({zip_size:.1f} MB)")
    return zip_path


def build_executable(include_cuda: bool = True,
                     create_zip_dist: bool = False,
                     output_base: Path = None):
    """Build the executable using PyInstaller (onedir mode)."""
    print("\nBuilding Transcribair executable...")
    print("-" * 50)

    # Use a build location outside OneDrive to avoid file-lock issues
    if output_base is None:
        output_base = Path(os.environ.get("TEMP", ".")) / "transcribair_build"

    dist_dir = output_base / "dist"
    build_dir = output_base / "build"

    for d in [dist_dir, build_dir]:
        if d.exists():
            print(f"Cleaning previous {d.name} directory...")
            for attempt in range(3):
                try:
                    shutil.rmtree(d)
                    break
                except PermissionError:
                    if attempt < 2:
                        print(f"  [WARNING] {d.name}/ locked, retrying in 3s...")
                        time.sleep(3)
                    else:
                        print(f"  [WARNING] Could not fully remove {d.name}/, "
                              "continuing anyway")

    dist_dir.mkdir(parents=True, exist_ok=True)
    build_dir.mkdir(parents=True, exist_ok=True)

    # Run PyInstaller
    print("\nRunning PyInstaller (onedir mode)...")
    print(f"  Output: {dist_dir}")
    print("  Excluding unused packages (matplotlib, scipy, pandas, etc.)")
    print("  UPX compression disabled (for compatibility)")
    print("  Debug symbols preserved (for stability)")
    print("")

    result = subprocess.run(
        [sys.executable, "-m", "PyInstaller", "build.spec", "--clean", "-y",
         "--distpath", str(dist_dir), "--workpath", str(build_dir)],
        capture_output=False
    )

    if result.returncode != 0:
        print("\n[ERROR] Build failed!")
        return False

    output_dir = dist_dir / "Transcribair"
    exe_path = output_dir / "Transcribair.exe"

    if not exe_path.exists():
        print(f"\n[ERROR] Expected output not found: {exe_path}")
        return False

    # Copy CUDA DLLs if requested
    cuda_count = 0
    if include_cuda:
        cuda_count = copy_cuda_dlls(output_dir)
    else:
        print("\n[INFO] Skipping CUDA DLLs (--no-cuda)")

    total_size = get_size_mb(output_dir)

    print("\n" + "=" * 50)
    print("[OK] Build successful!")
    print("=" * 50)
    print(f"\nOutput folder: {output_dir}")
    print(f"Executable:    {exe_path}")
    print(f"Total size:    {total_size:.1f} MB")
    print(f"\nBuild details:")
    print("  [OK] Excluded unused packages")
    if cuda_count > 0:
        print(f"  [OK] CUDA DLLs bundled ({cuda_count} files)")
    elif include_cuda:
        print("  [WARNING] No CUDA DLLs found (GPU will need system CUDA)")
    else:
        print("  [INFO] CPU-only build (no CUDA DLLs)")
    print("\nNote: First run will download the selected Whisper model.")
    print("      Models are cached in: %USERPROFILE%\\.transcribair\\models")

    # Create ZIP distribution if requested
    if create_zip_dist:
        create_zip(output_dir, include_cuda and cuda_count > 0)

    return True


def main():
    """Main build process."""
    parser = argparse.ArgumentParser(description="Build Transcribair executable")
    parser.add_argument("--no-cuda", action="store_true",
                        help="Build without CUDA DLLs (CPU-only, smaller)")
    parser.add_argument("--zip", action="store_true",
                        help="Create a distribution ZIP after building")
    args = parser.parse_args()

    print("Transcribair Build Script")
    print("=" * 50)
    if args.no_cuda:
        print("Mode: CPU-only (no CUDA DLLs)")
    else:
        print("Mode: CUDA-enabled (will bundle GPU DLLs if available)")

    if not check_requirements():
        print("\nPlease install missing requirements and try again.")
        sys.exit(1)

    # Build executable
    if build_executable(include_cuda=not args.no_cuda,
                        create_zip_dist=args.zip):
        print("\nBuild complete!")
    else:
        print("\nBuild failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()