import os
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
from pydoc import doc

# Configuration
ROOT_DIR = "./os"
MAX_WORKERS = 8  # Maximum parallel jobs for build docker image and configure project
MAX_WORKERS_COMPILE = 2 # number of compilation jobs t run

def find_dockerfiles(root_dir):
    """Recursively find all Dockerfiles in the given directory."""
    dockerfiles = []
    for dirpath, _, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename == "Dockerfile":
                dockerfiles.append(os.path.join(dirpath, filename))
    return dockerfiles


def tags_from_path(dockerfile):
    dir_path = os.path.dirname(dockerfile)
    tag = dir_path.replace(ROOT_DIR + "/", "").replace(os.path.sep, "_")
    return dir_path, tag


def build_docker_image(dockerfile):
    """Build a Docker image for the given Dockerfile."""
    dir_path, tag = tags_from_path(dockerfile)
    print(f"Starting build image step for {tag} (directory: {dir_path})...")
    try:
        log_file = f"_tmp/logs/{tag}/build_docker_image.log"
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        with open(log_file, "w") as log:
            result = subprocess.run(
                ["docker", "build", "-t", tag, dir_path],
                stdout=log,
                stderr=subprocess.STDOUT,
                check=True
            )
        return f"Successfully built {tag}"
    except subprocess.CalledProcessError:
        print(f"Failed to build {tag}. Check {log_file} for details.")
        return f"Failed to build {tag}"


in_image_source_dir = "/home/rublon-ssh"
in_image_build_dir = "/home/build"

def configure_rublon(dockerfile):
    """Build a Docker image for the given Dockerfile."""
    dir_path, tag = tags_from_path(dockerfile)
    print(f"Starting configure step for {tag} (directory: {dir_path})...")
    try:
        log_file = f"_tmp/logs/{tag}/configure_rublon_image.log"
        build_dir = f"_tmp/builds/{tag}"
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        os.makedirs(os.path.dirname(build_dir), exist_ok=True)
        with open(log_file, "w") as log:
            result = subprocess.run(
                # $tag /bin/bash /home/rublon-ssh/configure.sh $tag
                ["docker", "run",
                        "-v", f"./:{in_image_source_dir}",
                        "-v", f"./{build_dir}:{in_image_build_dir}",
                        tag,
                        f"{in_image_source_dir}/configure.sh", in_image_source_dir, in_image_build_dir
                ],
                stdout=log,
                stderr=subprocess.STDOUT,
                check=True
            )
        return f"Successfully built {tag}"
    except subprocess.CalledProcessError:
        print(f"Failed to build {tag}. Check {log_file} for details.")
        return f"Failed to build {tag}"


def build_rublon(dockerfile):
    """Build a Docker image for the given Dockerfile."""
    dir_path, tag = tags_from_path(dockerfile)
    print(f"Starting build for {tag} (directory: {dir_path})...")
    try:
        log_file = f"_tmp/logs/{tag}/compile_rublon_image.log"
        build_dir = f"_tmp/builds/{tag}"
        packages_dir = "_tmp/packages/"
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        os.makedirs(os.path.dirname(build_dir), exist_ok=True)
        os.makedirs(os.path.dirname(packages_dir), exist_ok=True)
        with open(log_file, "w") as log:
            result = subprocess.run(
                # $tag /bin/bash /home/rublon-ssh/configure.sh $tag
                ["docker", "run",
                        "-v", f"./:{in_image_source_dir}",
                        "-v", f"./{build_dir}:{in_image_build_dir}",
                        tag,
                        f"{in_image_source_dir}/build.sh", in_image_source_dir, in_image_build_dir, tag
                ],
                stdout=log,
                stderr=subprocess.STDOUT,
                check=True
            )
        return f"Successfully built {tag}"
    except subprocess.CalledProcessError:
        print(f"Failed to build {tag}. Check {log_file} for details.")
        return f"Failed to build {tag}"


def run_parralel(num, job, name, dockerfiles):
    # Use ThreadPoolExecutor for parallel builds
    with ThreadPoolExecutor(max_workers=num) as executor:
        future_to_dockerfile = {executor.submit(job, df): df for df in dockerfiles}
        for future in as_completed(future_to_dockerfile):
            dockerfile = future_to_dockerfile[future]
            try:
                result = future.result()
                print(result)
            except Exception as e:
                print(f"Job {name} failed for {dockerfile}: {e}")


def main():
    # Find all Dockerfiles
    dockerfiles = find_dockerfiles(ROOT_DIR)
    if not dockerfiles:
        print("No Dockerfiles found.")
        return

    run_parralel(MAX_WORKERS, build_docker_image, "Docker Image Build", dockerfiles)
    run_parralel(MAX_WORKERS, configure_rublon, "Configuring project", dockerfiles)
    run_parralel(MAX_WORKERS_COMPILE, build_rublon, "Building project", dockerfiles)


    print("All Docker builds completed.")

if __name__ == "__main__":
    main()