"""
Purpose: Automated construction of Docker images for repositories using profile registry.

Usage: python -m swesmith.build_repo.create_images --max-workers 4 -p django
"""

import argparse
import docker
import traceback
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

from swesmith.constants import Architecture
from swesmith.profiles import registry


def build_profile_image(profile, push=False):
    """
    Build a Docker image for a specific profile.

    Args:
        profile: A RepoProfile instance

    Returns:
        tuple: (profile_name, success: bool, error_message: str)
    """
    try:
        profile.create_mirror()
        profile.build_image()
        if push:
            profile.push_image()
        return (profile.image_name, True, None)
    except Exception as e:
        error_msg = f"Error building {profile.image_name}: {str(e)}"
        return (profile.image_name, False, error_msg)


def build_all_images(
    workers=4,
    repo_filter=None,
    proceed=False,
    push=False,
    force=False,
    arch=None,
):
    """
    Build Docker images for all registered profiles in parallel.

    Args:
        workers: Maximum number of parallel workers
        repo_filter: Optional list of repository name patterns to filter by (fuzzy matching)
        proceed: Whether to proceed without confirmation
        force: Force rebuild even if image already exists
        arch: Architecture to build for (e.g. "x86_64", "arm64")

    Returns:
        tuple: (successful_builds, failed_builds)
    """
    # Get all available profiles
    all_profiles = registry.values()

    # Update profile architecture if specified
    if arch:
        target_arch = arch
        print(f"Forcing build for architecture: {target_arch}")
        for profile in all_profiles:
            profile.arch = target_arch

    # Remove environments that have already been built
    client = docker.from_env()

    # Filter out profiles that already have images built (unless force is enabled)
    profiles_to_build = []
    if not force:
        for profile in all_profiles:
            try:
                # Check if image already exists
                client.images.get(profile.image_name)
            except docker.errors.ImageNotFound:
                profiles_to_build.append(profile)
    else:
        profiles_to_build = list(all_profiles)

    # Filter profiles if specified (fuzzy matching)
    if repo_filter:
        filtered_profiles = []
        for profile in profiles_to_build:
            # Check if any of the filter patterns appear in the image name
            if any(
                pattern.lower() in profile.image_name.lower() for pattern in repo_filter
            ):
                filtered_profiles.append(profile)
        profiles_to_build = filtered_profiles

    if not profiles_to_build:
        print("No profiles to build.")
        return [], []

    # Deduplicate profiles_to_build by image_name (more efficiently)
    profiles_to_build = list(
        OrderedDict(
            (profile.image_name, profile) for profile in profiles_to_build
        ).values()
    )

    print("Profiles to build:")
    for profile in sorted(profiles_to_build, key=lambda p: p.image_name):
        print(f"- {profile.image_name}")

    if not proceed:
        proceed = (
            input(
                f"Proceed with building {len(profiles_to_build)} images? (y/n): "
            ).lower()
            == "y"
        )
    if not proceed:
        return [], []

    # Build images in parallel
    successful, failed = [], []

    with tqdm(
        total=len(profiles_to_build), smoothing=0, desc="Building environment images"
    ) as pbar:
        with ThreadPoolExecutor(max_workers=workers) as executor:
            # Submit all build tasks
            future_to_profile = {
                executor.submit(build_profile_image, profile, push): profile
                for profile in profiles_to_build
            }

            # Process completed tasks
            for future in as_completed(future_to_profile):
                pbar.update(1)
                profile_name, success, error_msg = future.result()

                if success:
                    successful.append(profile_name)
                else:
                    failed.append(profile_name)
                    if error_msg:
                        print(f"\n{error_msg}")
                        traceback.print_exc()

    # Show results
    if len(failed) == 0:
        print("All environment images built successfully.")
    else:
        print(f"{len(failed)} environment images failed to build.")

    return successful, failed


def main():
    parser = argparse.ArgumentParser(
        description="Build Docker images for all registered repository profiles"
    )
    parser.add_argument(
        "-w",
        "--workers",
        type=int,
        default=4,
        help="Maximum number of parallel workers (default: 4)",
    )
    parser.add_argument(
        "-r",
        "--repos",
        type=str,
        nargs="+",
        help="Repository name patterns to build (fuzzy match, space-separated)",
    )
    parser.add_argument(
        "-y", "--proceed", action="store_true", help="Proceed without confirmation"
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Force rebuild even if image already exists",
    )
    parser.add_argument(
        "-p",
        "--push",
        action="store_true",
        help="Push built images to Docker Hub after building (default: False)",
    )
    parser.add_argument(
        "--list-envs", action="store_true", help="List all available profiles and exit"
    )
    parser.add_argument(
        "--arch",
        choices=[a.value for a in Architecture],
        help="Force build for specific architecture",
    )

    args = parser.parse_args()

    if args.list_envs:
        print("All execution environment Docker images:")
        for profile in registry.values():
            print(f"  {profile.image_name}")
        return

    successful, failed = build_all_images(
        workers=args.workers,
        repo_filter=args.repos,
        proceed=args.proceed,
        push=args.push,
        force=args.force,
        arch=args.arch,
    )

    if failed:
        print(f"- Failed builds: {failed}")
    if successful:
        print(f"- Successful builds: {len(successful)}")


if __name__ == "__main__":
    main()
