#!/usr/bin/python3
import os
import subprocess
import time
import invoke


@invoke.task(
    help={
        "mpsd_spack_ver": "MPSD Software stack version (defaults to current branch name)",
        "toolchain_list": "Comma separated list of toolchains to build (defaults to all toolchains)",
        "toolchain_base_dir": "Base directory to install the toolchains (defaults to '/opt_mpsd/')",
        "skip_build_cache": "Skip building the spack cache (defaults to False)",
        "skip_dir_check": "Skip checking if toolchains exists ( usefull when restarting after an error)",
    }
)
def build_toolchains(
    c,
    mpsd_spack_ver=None,
    toolchain_list=None,
    toolchain_base_dir="/opt_mpsd/",
    skip_build_cache=False,
    skip_dir_check=False,
):
    """
    Build toolchains using Spack.

    This task 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.
    """
    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_dir = os.getcwd()

    # Checks

    ## Check if current Git branch name and MPSD_SPACK_VER( If passed ) match
    if mpsd_spack_ver is None:
        mpsd_spack_ver = current_branch
    elif current_branch != mpsd_spack_ver:
        # Ensure that the current Git branch name matches MPSD_SPACK_VER
        raise invoke.Exit(
            f"Error: Current Git branch name {current_branch} \n\
            does not match argument for MPSD_SPACK_VER: {mpsd_spack_ver}"
        )

    ## Check if toolchains directory exists
    toolchains_path = os.path.join(
        toolchain_base_dir, mpsd_os, mpsd_spack_ver, mpsd_microarch
    )
    if not os.path.exists(toolchains_path):
        os.makedirs(toolchains_path)
    else:
        if not skip_dir_check:
            raise invoke.Exit(
                f"Error: Toolchains directory {toolchains_path} already exists. \n\
                Please remove it and try again."
            )
    ## 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 invoke.Exit(
                    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 c.cd(toolchains_path):
            c.run(f"echo '>>>> Building {toolchain}...' | tee -a {log_file}")
            c.run(
                f" bash {current_dir}/spack_setup.sh {flags} {toolchain}|tee -a {log_file} 2>&1"
            )
            # copy the octopus configs to the toolchain directory
            # c.run(f"cp -r {current_dir}/octopus {toolchains_path}")