cli.py

import argparse
import os
import tempfile
import git
import shutil
import stat
from app import (
    run_code2prompt,
    initialize_gemini_model,
    generate_dockerfile,
)

def clone_github_repo(repo_url, temp_dir):
    """Clones a GitHub repository to a temporary directory."""
    try:
        git.Repo.clone_from(repo_url, temp_dir)
        return temp_dir
    except git.exc.GitCommandError as e:
        print(f"Error cloning repository: {e}")
        return None

def on_rm_error(func, path, exc_info):
    """Handles errors when removing files or directories."""
    # Change file permissions and retry
    os.chmod(path, stat.S_IWRITE)
    os.unlink(path)

def remove_temp_dir(temp_dir):
    """Removes a temporary directory, handling permission errors."""
    try:
        shutil.rmtree(temp_dir, onerror=on_rm_error)
    except Exception as e:
        print(f"Warning: Could not delete temporary directory {temp_dir}: {e}")

def main():
    parser = argparse.ArgumentParser(description="Generate Dockerfiles from codebases using code2prompt and Gemini.")
    parser.add_argument("codebase_path", help="Path to the codebase (local directory or GitHub URL)")
    parser.add_argument("-i", "--include", help="Include patterns (e.g., *.py,*.js)")
    parser.add_argument("-e", "--exclude", help="Exclude patterns (e.g., *.txt,*.md)")
    parser.add_argument("-t", "--template", help="Path to a custom code2prompt template")
    parser.add_argument("-a", "--instructions", help="Additional instructions for Dockerfile generation")
    parser.add_argument("-o", "--output", default="Dockerfile", help="Output file name (default: Dockerfile)")
    args = parser.parse_args()

    # Determine if codebase_path is a URL or local path
    if args.codebase_path.startswith("http"):
        # Clone GitHub repository to a temporary directory
        temp_dir = tempfile.mkdtemp()
        print(f"Cloning repository {args.codebase_path} to {temp_dir}...")
        if clone_github_repo(args.codebase_path, temp_dir):
            codebase_path = temp_dir
        else:
            print(f"Error cloning repository: {args.codebase_path}")
            exit(1)
    else:
        codebase_path = args.codebase_path

    # Run code2prompt
    print("Running code2prompt...")
    code2prompt_output = run_code2prompt(
        codebase_path=codebase_path,
        include_patterns=args.include,
        exclude_patterns=args.exclude,
        template_path=args.template,
        json_output=False  # Assuming you want the raw prompt text for CLI
    )

    if code2prompt_output:
        # Initialize Gemini model
        print("Initializing Gemini model...")
        model = initialize_gemini_model()

        # Generate Dockerfile
        print("Generating Dockerfile...")
        dockerfile = generate_dockerfile(
            model,
            code2prompt_output,
            args.instructions
        )

        if dockerfile:
            # Write Dockerfile to output file
            with open(args.output, "w") as f:
                f.write(dockerfile)
            print(f"Dockerfile written to {args.output}")
        else:
            print("Error generating Dockerfile.")
    else:
        print("Error running code2prompt.")

    # Clean up temporary directory if it was created
    if args.codebase_path.startswith("http"):
        remove_temp_dir(temp_dir)

if __name__ == "__main__":
    main()