import os
from lib.shell import get_from_command


def get_compiler_version(compiler):
    if compiler == "gnu":
        cc = "gcc"
    elif compiler == "intel":
        cc = "icc"
    elif compiler == "aocc":
        cc = "clang"
    compstr = get_from_command([cc, "-dumpversion"]).strip()
    return compstr


def get_mpi_version(mpi):
    if mpi == "hpcx":
        hpcx_dir = os.getenv('HPCX_DIR')
        ver_file = hpcx_dir + "/" + "VERSION"
        with open(ver_file) as f:
            lines = f.readlines()
        mpistr = lines[0].split()[1].strip()
    elif mpi == "openmpi":
        rawstr = get_from_command(["ompi_info"]).splitlines()[1]
        mpistr = rawstr[rawstr.find(":")+2:]
    elif mpi == "intelmpi":
        rawstr = get_from_command(["ompi_info"]).splitlines()[0]
        mpistr = rawstr[rawstr.find("Version")+8:rawstr.find("Build")-1]
    elif mpi == "mpich":
        rawstr = get_from_command(["mpirun", "--version"])
        mpistr = rawstr.split()[4]
    return mpistr


def set_toolchain(compiler, mpi):
    if compiler == "gnu":
        if mpi == "hpcx" or mpi == "openmpi" or mpi == "mpich":
            cc = "mpicc"
            cxx = "mpic++"
            fc = "mpifort"
        elif mpi == "intel":
            cc = "mpicc"
            cxx = "mpicxx"
            fc = "mpif90"
    elif compiler == "aocc":
        if mpi == "hpcx" or mpi == "openmpi":
            cc = "mpicc"
            cxx = "mpic++"
            fc = "mpifort"
            os.environ["OMPI_MPICC"] = "clang"
            os.environ["OMPI_MPICXX"] = "clang++"
            os.environ["OMPI_MPIFC"] = "flang"
            os.environ["OMPI_MPIF90"] = "flang"
            os.environ["OMPI_MPIF77"] = "flang"
        elif (mpi == "intel"):
            cc = "\"mpicc -cc=clang\""
            cxx = "\"mpicxx -cxx=clang++\""
            fc = "\"mpif90 -fc=flang\""
    elif compiler == "intel":
        if mpi == "hpcx" or mpi == "openmpi":
            cc = "mpicc"
            cxx = "mpic++"
            fc = "mpifort"
            os.environ["OMPI_MPICC"] = "icc"
            os.environ["OMPI_MPICXX"] = "icpc"
            os.environ["OMPI_MPIFC"] = "ifort"
            os.environ["OMPI_MPIF90"] = "ifort"
            os.environ["OMPI_MPIF77"] = "ifort"
        elif mpi == "intel":
            cc = "mpiicc"
            cxx = "mpiicpc"
            fc = "mpiifort"
 
    # set environment variables
    os.environ["CC"] = cc
    os.environ["CXX"] = cxx
    os.environ["FC"] = fc
    os.environ["F90"] = fc
    os.environ["F77"] = fc

    return cc, cxx, fc