try:
    from rich import print
    from rich.prompt import Prompt

except ImportError:
    import sys

    print("You need to have `rich` installed to be able to run this script.")

    sys.exit(1)


import os


def list_folders(path: str) -> list[str]:
    """Lists all (immediate) subdirectories in a directory.
    
    Args:
        path (str): The path to list the subdirectories for.

    Returns:
        list[str]: A list of paths containing all the subdirectories for the specified path.
    """

    return [f.path.rsplit("/", 1)[-1] for f in os.scandir(path) if f.is_dir()]


variations = list_folders("./dockerfiles")
variation = Prompt.ask("What variation do you want to build?", choices=variations) if len(variations) > 1 else variations[0]

tags = list_folders(f"./dockerfiles/{variation}")
tag = Prompt.ask("What Python tag do you want to build?", choices=tags, default="all") if len(tags) > 1 else tags[0]


def build_image(variation: str, tag: str) -> bool:
    image = f"registry.nyeki.dev/docker/python/{variation}:{tag}"

    print(f"[yellow]Building and pushing [bold]{image}[/].[/]")

    commands = [
        f"cd ./dockerfiles/{variation}/{tag}",
        f"docker build -t {image} .",
        f"docker push {image}"
    ]

    if os.system(" && ".join(commands)) == 0:
        print(f"[green]Successfully built and pushed![/]")
        return True

    else:
        print(f"[red]The image couldn't be built or pushed.[/]")
        return False


if tag == "all":
    successes = 0

    for tag in tags:
        if build_image(variation, tag):
            successes += 1

    if successes == len(tags):
        print(f"[green]All tags were successfully built and pushed.[/]")

    elif successes != 0:
        print(f"[yellow]Only {successes}/{len(tags)} tags could be built and pushed.[/]")

    else:
        print(f"[red]No tag could be built nor pushed.[/]")

else:
    build_image(variation, tag)