#!/usr/bin/env python3
import os
import subprocess
import time
import argparse
import sys

about_tool="""
Build toolchains using Spack.

This function builds toolchains for toolchains at the appropriate directory for the current system architecture
and MPSD software stack version.
The toolchains are built using the bash script spack_setup.sh, and the results are logged.
"""
# Helper class to change directory via context manager
class os_chdir:
    def __init__(self, new_dir):
        self.new_dir = new_dir
        self.saved_dir = os.getcwd()
    
    def __enter__(self):
        os.chdir(self.new_dir)
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        os.chdir(self.saved_dir)


def build_toolchains(
    mpsd_spack_ver=None,
    toolchain_list=None,
    toolchain_base_dir="/opt_mpsd/",
    skip_build_cache=False,
    skip_dir_check=False,
):

    mpsd_os = os.environ["MPSD_OS"]
    mpsd_microarch = os.environ["MPSD_MICROARCH"]
    current_branch = (
        subprocess.run(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE
        )
        .stdout.decode()
        .strip()
    )
    current_commit_hash = (
        subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE
        )
    )
    current_dir = os.getcwd()

    # Checks

    ## Check that some mpsd_spack_ver is passed.
    if mpsd_spack_ver is None:
        raise ValueError(
            "Error: MPSD_SPACK_VER not passed. Please pass the MPSD Software stack version."
        )
    ## Check if toolchains directory exists
    spack_env_path = os.path.join(toolchain_base_dir, mpsd_os, mpsd_spack_ver)
    toolchains_path = os.path.join(spack_env_path, mpsd_microarch)
    if not os.path.exists(toolchains_path):
        os.makedirs(toolchains_path)
    else:
        if not skip_dir_check:
            raise ValueError(
                f"Error: Toolchains directory {toolchains_path} already exists. \n\
                Please remove it and try again."
            )
    ### Clone spack-env repo at spack_env_path
    with os.chdir(spack_env_path + "/spack-environments"):
        subprocess.run("git clone git@gitlab.gwdg.de:mpsd-cs/spack-environments.git .", shell=True, check=True)
        subprocess.run(f"git checkout {mpsd_spack_ver}", shell=True, check=True)
    ## Check if TOOLCHAIN_LIST is valid
    available_toolchains = os.listdir("toolchains")

    if skip_build_cache:
        flags = "-b"
    else:
        flags = ""

    if toolchain_list is None:
        toolchains = available_toolchains
    else:
        toolchains = toolchain_list.split(",")
        # if not set(toolchains)<=set(available_toolchains):
        for toolchain in toolchains:
            if toolchain not in available_toolchains:
                raise ValueError(
                    f"Error: Toolchain '{toolchain}' not found in toolchains directory. \n\
                    Please check the toolchain argument and try again."
                )

    # Build toolchains

    log_file = f"build_toolchains_{mpsd_spack_ver}_{time.strftime('%Y%m%d-%H%M%S')}.log"
    print(f"Building at {toolchains_path}...")
    for toolchain in toolchains:
        with os_chdir(toolchains_path):
            print(f">>>> Building {toolchain}...")
            subprocess.run(
                f" bash {spack_env_path}/spack-environments/spack_setup.sh {flags} {toolchain}|tee -a {log_file} 2>&1", 
                shell=True, check=True
            )
            # copy the octopus configs to the toolchain directory
            # subprocess.run(f"cp -r {current_dir}/octopus

def prepare_env(release, target_dir, skip_dir_check):
    if target_dir == "DEFAULT":
        target_dir = f"/opt_mpsd/{os.environ['MPSD_OS']}/{release}/{os.environ['MPSD_MICROARCH']}"
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    else:
        if not skip_dir_check:
            raise ValueError(
                f"Error: Target directory {target_dir} already exists. \n\
                Please remove it and try again."
            )
    with os.chdir(target_dir):
        subprocess.run("git clone https://gitlab.gwdg.de/mpsd-cs/spack-environments.git")
        subprocess.run(f"git checkout {release}")
def install_toolchains(release, install, target_dir, force_reinstall, skip_build_cache, skip_dir_check):
    prepare_env(release, target_dir, skip_dir_check)

def start_new_env(set_up,from_release):
    pass

def remove_toolchain(release, remove):
    pass
def main():
    parser = argparse.ArgumentParser(description=about_tool)
    parser.add_argument('--release', type=str,
                        help='Specify the release version to install')
    parser.add_argument('--target-directory', type=str,
                        help='Specify the target directory for installation (use DEFAULT to use /opt_mpsd/<MPSD_OS>/<MPSD_RELEASE>/<MPSD_MICROARCH)')
    parser.add_argument('--install', nargs='+',
                        help='Specify toolchain(s) to install eg foss2021a-mpi (use ALL to install all available toolchains)')
    parser.add_argument('--remove', type=str,
                        help='Specify toolchain to remove')
    parser.add_argument('--set-up', type=str,
                        help='Start a new software environment version, must specify --from <release>')
    parser.add_argument('--from', dest='from_release', type=str,
                        help='Specify the release version to start from')
    parser.add_argument('--force-reinstall', action='store_true',
                        help='Delete and reinstall an existing toolchain directory')
    parser.add_argument('--skip-build-cache', action='store_true',
                        help='Skip Spack build cache during installation')
    parser.add_argument('--skip-dir-check', action='store_true',
                        help='Skip checking if the target directory already exists')

    args = parser.parse_args()

    if args.release is None:
        parser.print_help()
        sys.exit(1)

    if args.remove:
        remove_toolchain(args.release, args.remove)
    elif args.set_up:
        start_new_env(args.set_up, args.from_release)
    else:
        target_dir = args.target_directory

        if args.install:
            install_toolchains(args.release, args.install, target_dir, args.force_reinstall,
                               args.skip_build_cache, args.skip_dir_check)
        else:
            prepare_env(args.skip_dir_check, args.release, target_dir)
    

if __name__ == "__main__":
    main()