#!/usr/bin/python3
"""
Compiler ID  handling for dhfortran

Copyright (C) 2025 Alastair McKinstry <mckinstry@debian.org>
Released under the GPL-3 GNU Public License.

This module includes all compiler-specific stuff, typically passed to
"""

import os
from pathlib import Path
from subprocess import check_output
from shutil import which


multiarch = (
    check_output(["dpkg-architecture", "-qDEB_HOST_MULTIARCH"]).strip().decode("utf-8")
)
libdir = f"/usr/lib/{multiarch}"
flibdir = f"{libdir}/fortran"

# To be updated as gfortran, flang etc are updated  for new versions
default_compilers = {
    "gfortran": "gfortran-15",
    "flang": "flang19",
    "flangext": "flang-to-external-fc-18",
}

all_compilers = {
    "gfortran": ["gfortran-7", "gfortran-13", "gfortran-14", "gfortran-15"],
    "flang": ["flang-new-18", "flang-new-19", "flang-new-20", "flang-new-21"],
}

# Compiler information.
# ABI names compatible with CMake Identifiers (lower case)
# first exe name is the default

compilers = {
    "gfortran-7": {"exe": ["gfortran-7"], "abi": "GNU", "mod": "gfortran-mod-14"},
    "gfortran-13": {"exe": ["gfortran-13"], "abi": "GNU", "mod": "gfortran-mod-15"},
    "gfortran-14": {"exe": ["gfortran-14"], "abi": "GNU", "mod": "gfortran-mod-15"},
    "gfortran-15": {"exe": ["gfortran-15"], "abi": "GNU", "mod": "gfortran-mod-16"},
    "flang-new-17": {"exe": ["flang-new-17"], "abi": "Flang", "mod": "flang-mod-1"},
    "flang-new-18": {"exe": ["flang-new-18"], "abi": "Flang", "mod": "flang-mod-1"},
    "flang-new-18": {"exe": ["flang-new-19"], "abi": "flang", "mod": "flang-mod-1"},
    "flang-new-20": {"exe": ["flang-new-20"], "abi": "flang", "mod": "flang-mod-1"},
    "flang-21": {
        "exe": ["flang-new-21", "flang-21"],
        "abi": "flang",
        "mod": "flang-mod-1",
    },
    "lfortran": {"exe": ["lfortran"], "abi": "lfortran", "mod": "lfortran-mod-0"},
    "flang-to-external-fc-17": {
        "exe": ["flang-to-external-fc-17"],
        "abi": "flangext",
        "mod": "flangext-mod-15",
    },  # flangext not in CMake list. Now obsolete anyway
    "flang-to-external-fc-18": {
        "exe": ["flang-to-external-fc-18"],
        "abi": "flangext",
        "mod": "flangext-mod-15",
    },
    # Commercial compilers. Need to be checked TODO
    "intel": {"exe": ["ifx", "ifc", "efc", "iort"], "abi": "Intel", "mod": "UNKNOWN"},
    "pgi": {
        "exe": ["pgfortran", "pgf95", "pgf90", "pgf77"],
        "abi": "pgi",
        "mod": "UNKNOWN",
    },
}

# Compatible ABIs. TODO
compatible = {
    "GNU": ["GNU"],
}


# See  #957692
# TODO: Rework this with  Guillem, a plugin mechanism for dpkg-buildflags

fc_flags_append = {
    "gfortran-10": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
    "gfortran-11": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
    "gfortran-12": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
    "gfortran-13": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
    "gfortran-14": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
    "gfortran-15": ["-fallow-invalid-boz", "-fallow-argument-mismatch"],
}

fc_flags_strip = {
    "flang-7": ["-g"],
    "flang-new-18": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
    ],
    "flang-new-19": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
    ],
    "flang-new-20": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
    ],
    "flang-21": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
    ],
}


def get_fc_flavor(fc: str) -> str:
    """Given a compiler name, work out flavor"""
    # TODO: currently assumes fc is a full path.
    # TODO: mpifort , h5fc ?
    # resolve symlinks and strip prefix for cross-compile , etc
    # assumes no broken symlinks ...
    base = Path(fc).resolve().name
    m = check_output(["dpkg-architecture", "-qDEB_HOST_GNU_TYPE"]).strip().decode("utf-8")
    if base.startswith(m):
        base = base[len(m) + 1 :]
    for f in compilers:
        if base in compilers[f]["exe"]:
            return f
    raise Exception(f"Can't recognize compiler {fc}")


def get_fmoddir(flavor: str) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    modv = compilers[f]["mod"]
    return f"{flibdir}/{modv}"


def get_flibdir(flavor: str) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    abi = compilers[f]["abi"].lower()
    return f"{flibdir}/{abi}"


def get_fc_flags(flavor: str) -> str:
    # TODO
    fcflags = check_output(["dpkg-buildflags", "--get", "FCFLAGS_FOR_BUILD"]).strip().decode("utf-8")
    # TODO match % in flags
    if flavor in fc_flags_strip:
        for r in fc_flags_strip[flavor]:
            if r in fcflags:
                fcflags.remove(r)
    if flavor in fc_flags_append:
        fcflags += " ".join(fc_flags_append[flavor])
    return fcflags


def get_flavor(flavor: str) -> str:
    # FC_FLAVOR matches fortran-support.mk loop-over-flavors variable
    # TODO merge get_fc_flavor

    if flavor is not None:
        return flavor
    if "FC_FLAVOR" in os.environ:
        return os.environ["FC_FLAVOR"]
    try:
        return get_fc_flavor(get_fc())
    except Exception as ex:
        raise Exception(
            "Either --flavor is set, or FC_FLAVOR or FC are defined in the environment"
        )


def get_preferred(preferred: str, flavor: str) -> str:
    """When links are set up and multiple versions of a library can/may be installed.
    which one is installed ? which links are set up?

    flavor: must not be None
    """
    if preferred is not None:
        return preferred
    if "PREFERRED" in os.environ:
        return os.environ["PREFERRED"]

    if abi := get_abi_vendor(flavor) in default_compilers:
        return default_compilers[abi].lowercase()
    else:
        # fallback to default
        return get_flavor(None)


def get_fc(fc=None):
    if fc is None and "FC" in os.environ:
        f = os.environ["FC"]
    else:
        f = fc if fc else "/etc/alternatives/f95"
    fullpath = which(f)
    if fullpath is None:
        raise Exception(f"fc compiler {f} is broken; bad symlink?")
    return Path(fullpath).resolve().name


def get_f77(f77=None):
    if f77 is None and "F77" in os.environ:
        f = os.environ["F77"]
    else:
        f = f77 if f77 else "/etc/alternatives/f77"
    fullpath = which(f)
    if fullpath is None:
        raise Exception(f"f77 compiler {f} is broken; bad symlink?")
    return Path(fullpath).resolve().name


def get_fc_default(fc=None) -> str:
    """Return the default Fortran compiler"""
    if "FC_DEFAULT" in os.environ:
        return os.environ["FC_DEFAULT"]
    x = get_fc(fc)
    return get_fc_flavor(x)


def get_fc_optional(fc=None) -> str:
    """Return the list of other compilers present on the system"""
    compilers_present = {}
    d = get_fc_default(fc)
    for flavor in compilers:
        for exes in compilers[flavor]["exe"]:
            if os.path.exists(f"/usr/bin/{exes}"):
                if flavor != d and flavor not in compilers_present:
                    compilers_present[flavor] = flavor
    return " ".join(compilers_present)


def get_pkg_config_path(flavor: str) -> str:
    pkg_config_path = (
        os.environ["PKG_CONFIG_PATH"] if "PKG_CONFIG_PATH" in os.environ else ""
    )
    return ":".join([flibdir + "/" + flavor + "/pkgconfig", pkg_config_path])


def get_cmake_path(flavor: str) -> str:
    cmake_module_path = (
        os.environ["CMAKE_MODULE_PATH"] if "CMAKE_MODULE_PATH" in os.environ else ""
    )
    return ":".join([flibdir + "/" + flavor + "/cmake", cmake_module_path])


# TODO: Drop?
def fc_target_arch_flang(fcflags: str) -> str | None:
    """Return target triple from flang ,
    or None for default
    """
    # TODO not used or tested yet
    flags = fcflags.split()

    if "-target" in flags:
        try:
            target = flags[flags.index("-target") + 1]
        except Exception as ex:
            raise Exception("Parse error finding -target in flang flags list")
        try:
            return triplet[target]
        except:
            raise Exception(f"Unknown arch target {target} in flang flags")
    return None


def get_env(flavor: str) -> str:
    # Used internally to set env flags for debhelper perl, but also debugging
    fc, f77 = get_fc(), get_f77()
    fc_default = get_fc_default(fc)
    fc_optional = get_fc_optional(fc)

    if flavor is None:
        flavor = fc_default
    return f"""FC={fc}
F77={f77}
FLAVOR={flavor}
FC_DEFAULT={fc_default}
FC_OPTIONAL={fc_optional}
PKG_CONFIG_PATH={get_pkg_config_path(fc_default)}
CMAKE_MODULE_PATH={get_cmake_path(fc_default)}
"""


def fc_target_arch_lfortran(fcflags: str) -> str:
    """Return the target triple for an lfortran compilation
    TODO: Needs more work
    """

    flags = fcflags.split()

    if "-target" in flags:
        try:
            target = flags[flags.index("-target") + 1]
        except Exception as ex:
            raise Exception("Parse error finding -target in lfortran flags list")

        # Translate target archs to GNU triplet
        archs = {
            "aarch64": "aarch64-linux-gnu",  #   - AArch64 (little endian)
            "aarch64_32": None,  #  - AArch64 (little endian ILP32) unsupported
            "aarch64_be": "aarch64_be-linux-gnu",  #  AArch64 (big endian) unsupported on Debian
            "arm64": "aarch64-linux-gnu",  # ARM64 (little endian)
            "arm64_32": None,  #   - ARM64 (little endian ILP32)
            "x86": "i386-linux-gnu",  #     - 32-bit X86: Pentium-Pro and above
            "x86-64": "x86_64-linux-gnu",  #     - 64-bit X86: EM64T and AMD64
        }
        return archs[target]
    return None


def get_fc_exe(flavor=None) -> str:
    """Return compiler name,  suitable for -DCMAKE_Fortran_COMPILER=$(call get_fc_exe,XXX)"""
    # TODO: cross-compilation ?

    f = flavor if flavor in compilers else get_fc_default()
    return compilers[f]["exe"][0]


def get_abi_vendor(flavor: str) -> str:
    """Return vendor name.
    Names compatible with CMake
    """

    # names from CMake 3.11 Fortran ID

    if flavor.startswith("gfortran") or flavor.startswith("flangext"):
        return "gnu"
    if flavor.startswith("flang"):
        # flang can mean _Fortran_COMPILER_NAMES_Flang or _Fortran_COMPILER_NAMES_LLVMFlang
        return "flang"
    if flavor.startswith("lfortran"):
        # Not presumed to be ABI-compatible with gfortran, but likely
        return "LCC"
    if flavor.startswith("path"):
        # pathf2003 pathf95 pathf90
        return "PathScale"
    if flavor.startswith("pgf"):
        # pgf95 pgfortran pgf90 pgf77
        return "PGI"
    if flavor.startswith("nvfortran"):
        return "NVHPC"
    if flavor.startswith("nagfor"):
        return "NAG"
    if flavor.startswith("xlf"):
        # set(_Fortran_COMPILER_NAMES_XL        xlf)
        # set(_Fortran_COMPILER_NAMES_VisualAge xlf95 xlf90 xlf)
        return "VisualAge"
    if flavor.startswith("af"):
        # af95 af90 af77
        return "Absoft"
    if flavor in ["ifort", "ifc", "efc", "ifx"]:
        return "Intel"

    raise Exception(f"ABI Vendor unknown for flavor {flavor}")


if __name__ == "__main__":
    import pytest

    pytest.main(["tests/compilers.py"])