#!/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 click.exceptions import UsageError
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")
)

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

all_compilers = {
    "gfortran": [
        "gfortran-7",
        "gfortran-13",
        "gfortran-14",
        "gfortran-15",
        "gfortran-16",
    ],
    "flang": ["flang-new-18", "flang-new-19", "flang-new-20", "flang-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",
        "compatible_moddirs": ["gfortran-mod-14", "flangext-mod-14"],
    },
    "gfortran-13": {
        "exe": ["gfortran-13"],
        "abi": "GNU",
        "mod": "gfortran-mod-15",
        "compatible_moddirs": ["gfortran-mod-15", "flangext-mod-15"],
    },
    "gfortran-14": {
        "exe": ["gfortran-14"],
        "abi": "GNU",
        "mod": "gfortran-mod-15",
        "compatible_moddirs": ["gfortran-mod-15", "flangext-mod-15"],
    },
    "gfortran-15": {
        "exe": ["gfortran-15"],
        "abi": "GNU",
        "mod": "gfortran-mod-16",
        "compatible_moddirs": ["gfortran-mod-16", "flangext-mod-16"],
    },
    "gfortran-16": {
        "exe": ["gfortran-16"],
        "abi": "GNU",
        "mod": "gfortran-mod-16",
        "compatible_moddirs": ["gfortran-mod-16", "flangext-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-19": {"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",
        "compatible_moddirs": ["flangext-mod-15", "gfortran-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",
        "compatible_moddirs": ["flangext-mod-15", "gfortran-mod-15"],
    },
    # Commercial compilers. Need to be checked TODO
    "intel": {
        "exe": ["ifx", "ifc", "efc", "iort"],
        "abi": "Intel",
        "mod": "UNKNOWN",
        "mod_arg": "-module ",
    },
    "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

# fPIC over the top and unnecessary in many cases, but workaround for fpm
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", "-fPIC" ],
    "gfortran-16": ["-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%",
        "-fcf-protection",
    ],
    "flang-new-19": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
        "-fcf-protection",
    ],
    "flang-new-20": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
        "-fcf-protection",
    ],
    "flang-21": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
        "-fcf-protection",
    ],
    "lfortran": [
        "-mbranch-protection=standard",
        "-fstack-protector-strong",
        "-fstack-clash-protection",
        "-ffile-prefix-map%",
        "-fcf-protection",
    ],
}

# Empty for now
ld_flags_append = {}

ld_flags_strip = {
    "flang-new-18": [
        "-mbranch-protection=standard",
    ],
    "flang-new-19": [
        "-mbranch-protection=standard",
    ],
    "flang-new-20": [
        "-mbranch-protection=standard",
    ],
    "flang-21": [
        "-mbranch-protection=standard",
    ],
}


def get_fc_flavor_arch(fc=None) -> str:
    """Given a compiler name, work out flavor and host arch
    Works with cross-compile, validates fc or fails with UsageError
    """
    # TODO: mpifort , h5fc ?
    # resolve symlinks and strip prefix for cross-compile , etc
    orig_fc = fc
    arch = multiarch
    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:
        fullpath = which(default_compilers['gfortran'])
    if fullpath is None:
        raise Exception(f"fc compiler {orig_fc} is broken; bad symlink?")
    fc = Path(fullpath).resolve().name
    if fc in default_compilers:
        fc = default_compilers[fc]
    m = (
        check_output(["dpkg-architecture", "-qDEB_BUILD_GNU_TYPE"])
        .strip()
        .decode("utf-8")
    )
    if fc.startswith(m):
        fc = fc[len(m) + 1 :]
    # cross-compiling
    m = (
        check_output(["dpkg-architecture", "-qDEB_HOST_GNU_TYPE"])
        .strip()
        .decode("utf-8")
    )
    if fc.startswith(m):
        fc = fc[len(m) + 1 :]
        arch = m
    for f in compilers:
        if fc in compilers[f]["exe"]:
            return f, arch
    raise UsageError(f"Can't recognize compiler {fc}")


def get_fmoddir(flavor: str, host_arch: str = None) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    host_arch = multiarch if host_arch is None else host_arch
    flibdir = f"/usr/lib/{host_arch}/fortran"
    modv = compilers[f]["mod"]
    return f"{flibdir}/{modv}"


def get_fmoddirs(flavor: str, host_arch: str) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    host_arch = multiarch if host_arch is None else host_arch
    flibdir = f"/usr/lib/{host_arch}/fortran"
    if "compatible_moddirs" in compilers[f]:
        return [f"{flibdir}/{modv}" for modv in compilers[f]["compatible_moddirs"]]
    else:
        return [get_fmoddir(f)]


def get_flibdir(flavor: str, host_arch: str) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    host_arch = multiarch if host_arch is None else host_arch
    flibdir = f"/usr/lib/{host_arch}/fortran"
    abi = compilers[f]["abi"].lower()
    return f"{flibdir}/{abi}"


def get_flibdirs(flavor: str, host_arch: str) -> str:
    f = flavor if flavor else default_compilers["gfortran"]
    host_arch = multiarch if host_arch is None else host_arch
    if "compatible_libdirs" in compilers[f]:
        flibdir = f"/usr/lib/{host_arch}/fortran"
        return ["f{flibdir}/{lib}" for lib in compilers[f]["compatible_libdirs"]]
    else:
        return [get_flibdir(f)]


def get_fc_flags(flavor: str, host_arch: str) -> str:
    fcflags = (
        check_output(["dpkg-buildflags", "--get", "FCFLAGS_FOR_BUILD"])
        .strip()
        .decode("utf-8")
    )
    fcflags_list = fcflags.split()
    # % in flag means wildcard match after that poin
    if flavor in fc_flags_strip:
        for r in fc_flags_strip[flavor]:
            if "%" in r:
                prefix = r[: r.index("%")]
                for x in fcflags_list:
                    if x.startswith(prefix):
                        fcflags_list.remove(x)
            else:
                if r in fcflags_list:
                    fcflags_list.remove(r)
    fcflags = " ".join(fcflags_list)
    if flavor in fc_flags_append:
        fcflags += " " + " ".join(fc_flags_append[flavor])
    fmoddir = get_fmoddir(flavor,host_arch)  
    fcflags += f" -I{fmoddir}"

    return fcflags


def get_ld_flags(flavor: str, arch: str) -> str:
    """Get the LDFLAGS to use if $FC is used as a linker.
    TODO: dpkg-buildflags presumes GNU linker, breaks when flang used
    TODO: multiarch aware
    """
    ldflags = (
        check_output(["dpkg-buildflags", "--get", "LDFLAGS_FOR_BUILD"])
        .strip()
        .decode("utf-8")
    )
    ldflags_list = ldflags.split()
    # % in flag means wildcard match after that poin
    if flavor in ld_flags_strip:
        for r in ld_flags_strip[flavor]:
            if "%" in r:
                prefix = r[: r.index("%")]
                for x in ldflags_list:
                    if x.startswith(prefix):
                        ldflags_list.remove(x)
            else:
                if r in ldflags_list:
                    ldflags_list.remove(r)
    ldflags = " ".join(ldflags_list)
    if flavor in ld_flags_append:
        ldflags += " " + " ".join(ld_flags_append[flavor])
    
    flibdir = get_flibdir(flavor,arch)  
    ldflags += f" -L{flibdir}"

    return ldflags


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"]
    abi = get_abi_vendor(flavor)
    if abi in default_compilers:
        return default_compilers[abi].lower()
    else:
        # fallback to default
        flavor, _ = get_fc_flavor_arch(None)
        return flavor


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:
        fullpath = which(default_compilers['gfortran'])
    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"]
    default, _ = get_fc_flavor_arch(fc)
    return default


def get_fc_optional(fc=None) -> set[str]:
    """Return the list of other compilers present on the system
    Remove the default from the list if given"""
    d = get_fc_default(fc)
    if "FC_OPTIONAL" in os.environ:
        flavors =  set(f for f in os.environ["FC_OPTIONAL"].split(" "))
        if d in flavors:
            flavors.remove(d)
        return flavors
    else:
        compilers_present = set()
        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.add(flavor)
        return compilers_present


def get_pkg_config_path(flavor: str, arch: str = None) -> str:
    arch = multiarch if arch is None else arch
    flibdir = f"/usr/lib/{arch}/fortran"
    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, arch: str = None) -> str:
    arch = multiarch if arch is None else arch
    flibdir = f"/usr/lib/{arch}/fortran"
    cmake_module_path = (
        os.environ["CMAKE_MODULE_PATH"] if "CMAKE_MODULE_PATH" in os.environ else ""
    )
    return ":".join([flibdir + "/" + flavor + "/cmake", cmake_module_path])


def get_env(flavor: str) -> str:
    # Used internally to set env flags for debhelper perl, but also debugging
    fc, arch = get_fc_flavor_arch(flavor)
    f77 = get_f77()
    fc_default = get_fc_default(fc)
    fc_optional = " ".join(get_fc_optional(fc))
    fc_flags = get_fc_flags(fc, arch)
    fc_ldflags = get_ld_flags(fc, arch)
    return f"""FC={fc}
F77={f77}
FC_DEFAULT={fc_default}
FC_OPTIONAL={fc_optional}
PKG_CONFIG_PATH={get_pkg_config_path(fc)}
CMAKE_MODULE_PATH={get_cmake_path(fc)}
FPM_FC={fc}
FPM_FFLAGS={fc_flags}
FPM_LDFLAGS={fc_ldflags}
"""


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

    flags = fcflags.split()

    if "-target" in flags:
        try:
            target = flags[flags.index("-target") + 1]
        except Exception:
            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"])
