main.py

# main.py
import concurrent.futures
import argparse
import sys
from email.mime import image
from pathlib import Path
from db import init_db, save_result, get_completed_packages, get_failed_packages
from recipe import collect_recipes, find_recipes_by_names, dump_top_k_packages, read_packages_from_file
from builder import build_in_docker, build_docker_image
from progress import build_with_progress, create_progress_bar, print_build_summary

def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description="Build Conan packages for multiple target architectures")
    
    parser.add_argument("--repo-path", type=Path, default="conan-center-index",
                       help="Path to conan-center-index repository (default: conan-center-index)")
    parser.add_argument("--logs-path", type=Path, default="logs",
                       help="Path to store build logs (default: logs)")
    parser.add_argument("--cache-path", type=Path, default="conan_cache",
                       help="Path to Conan cache directory (default: conan_cache)")
    parser.add_argument("--db-path", type=Path, default="build_results.sqlite",
                       help="Path to SQLite database file (default: build_results.sqlite)")
    parser.add_argument("--profile", type=Path, default="conan_profiles/riscv64-linux-gcc",
                       help="Path to Conan profile (default: conan_profiles/riscv64-linux-gcc)")
    parser.add_argument("--image", type=str, default="conan-builder",
                       help="Docker image name (default: conan-builder)")
    parser.add_argument("--platform", type=str, default="linux/riscv64",
                       help="Target platform for Docker builds (default: linux/riscv64)")
    parser.add_argument("--max-workers", type=int, default=2,
                       help="Maximum number of parallel build workers (default: 2)")
    parser.add_argument("--top-k", type=int, default=None, metavar="K",
                       help="Limit number of packages to process. Used for dumping top-k packages or limiting builds (no limit if not specified)")
    parser.add_argument("--packages", nargs='+', metavar="PACKAGE",
                       help="Build specific packages by name (overrides database check). Example: --packages zlib boost nlohmann_json")
    parser.add_argument("--package-file", type=Path, metavar="FILE",
                       help="Build packages listed in file (one package per line). Overrides database check.")
    parser.add_argument("--build-image", action="store_true",
                       help="Build the Docker image before starting package builds")
    parser.add_argument("--rebuild-failed", action="store_true",
                       help="Rebuild only packages that previously failed (from database)")
    parser.add_argument("--dump-top-k-to", type=Path, metavar="FILE",
                       help="Dump top-k packages (by commit count) to file and exit. Use --top-k to specify count (default: 100)")
    parser.add_argument("--force-rebuild", action="store_true",
                       help="Force rebuild packages even if already completed (applies to all modes)")
    
    return parser.parse_args()

def filter_completed_packages(recipes, conn, force_rebuild=False):
    """
    Filter out completed packages unless force rebuild is enabled
    
    Args:
        recipes: List of recipe dictionaries
        conn: Database connection
        force_rebuild: If True, don't filter completed packages
        
    Returns:
        list: Filtered recipes (excluding completed ones unless force_rebuild=True)
    """
    if force_rebuild:
        return recipes
    
    completed_packages = get_completed_packages(conn)
    filtered = [r for r in recipes if r['name'] not in completed_packages]
    
    skipped_count = len(recipes) - len(filtered)
    if skipped_count > 0:
        print(f"Skipping {skipped_count} already completed packages (use --force-rebuild to override)")
    
    return filtered

def main():
    args = parse_args()
    
    # Build Docker image if requested
    if args.build_image:
        if not build_docker_image(args.image, args.platform):
            print("Failed to build Docker image. Exiting.")
            sys.exit(1)
        print(f"Docker image '{args.image}' built successfully.")
        return

    # Get all recipes from conan center-index repo.
    recipes = collect_recipes(args.repo_path)
    print(f"Found {len(recipes)} recipes.")
    
    # Dump top-k packages if requested
    if args.dump_top_k_to:
        # Default to 100 if --top-k not specified for dump operation
        k = args.top_k if args.top_k is not None else 100
        dump_top_k_packages(recipes, args.dump_top_k_to, k)
        return

    # Initialize database and logs for build operations
    args.logs_path.mkdir(exist_ok=True)
    conn = init_db(args.db_path)
    conn.execute("PRAGMA journal_mode=WAL;")
    
    # Validate mutually exclusive package selection options
    package_options = [args.packages, args.package_file, args.rebuild_failed]
    if sum(opt is not None and opt is not False for opt in package_options) > 1:
        print("Error: --packages, --package-file, and --rebuild-failed are mutually exclusive")
        sys.exit(1)
    
    # Require one of the package selection modes
    if not any(opt is not None and opt is not False for opt in package_options):
        print("Error: Must specify one of --packages, --package-file, or --rebuild-failed")
        print("Use --help for usage information")
        sys.exit(1)

    if args.packages:
        # Specific packages mode - build only specified packages
        print(f"Building specific packages: {', '.join(args.packages)}")
        
        found_recipes, not_found = find_recipes_by_names(recipes, args.packages, "packages")
        pending = filter_completed_packages(found_recipes, conn, args.force_rebuild)
    
    elif args.package_file:
        # Package file mode - build packages listed in file
        package_names = read_packages_from_file(args.package_file)
        print(f"Building packages from file: {', '.join(package_names)}")
        
        found_recipes, not_found = find_recipes_by_names(recipes, package_names, "packages from file")
        pending = filter_completed_packages(found_recipes, conn, args.force_rebuild)
    
    elif args.rebuild_failed:
        # Rebuild failed packages mode - build only packages that previously failed
        failed_packages = get_failed_packages(conn)
        print(f"Found {len(failed_packages)} failed packages in database")
        
        if not failed_packages:
            print("No failed packages found in database. Nothing to rebuild!")
            conn.close()
            return
        
        print(f"Rebuilding failed packages: {', '.join(failed_packages)}")
        
        found_recipes, not_found = find_recipes_by_names(recipes, failed_packages, "failed packages")
        # For failed packages, we always want to rebuild them regardless of force_rebuild flag
        pending = found_recipes
        
        # Apply limit to failed packages if specified
        if args.top_k and len(pending) > args.top_k:
            print(f"Limiting rebuild to first {args.top_k} failed packages")
            pending = pending[:args.top_k]
    
    # Apply general limit if specified (for packages and package-file modes)
    if args.top_k and len(pending) > args.top_k:
        print(f"Limiting to first {args.top_k} packages")
        pending = pending[:args.top_k]
    
    print(f"{len(pending)} recipes pending build.")

    for r in pending:
        print(f"  {r['name']} ({r['commits']} commits, header-only: {r['header_only']})")

    if not pending:
        print("No packages to build!")
        conn.close()
        return

    print(f"\nStarting builds with {args.max_workers} workers...")
    
    # Create progress bar
    with create_progress_bar(len(pending), "Waiting for builds...") as pbar:
        completed_results = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
            futures = [
                executor.submit(
                    build_with_progress,
                    pbar,
                    build_in_docker,
                    recipe['name'],
                    recipe["version"],
                    recipe["path"],
                    args.profile,
                    args.image,
                    args.platform,
                    args.cache_path,
                    args.logs_path,
                )
                for recipe in pending
            ]
            for f in concurrent.futures.as_completed(futures):
                try:
                    result = f.result()
                    save_result(conn, result)
                    completed_results.append(result)
                except Exception as e:
                    print(f"\nError in build: {e}")
        
        # Final update
        pbar.set_description("Completed")
        
    # Print summary using the progress module
    print_build_summary(completed_results)
    conn.close()

if __name__ == "__main__":
    main()