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

Copyright (C) 2025-2026 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
import functools
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-21",
    "flangext": "flang-to-external-fc-18",
}


# 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",
    },
    "flang-22": {
        "exe": ["flang-new-22", "flang-22"],
        "abi": "flang",
        "mod": "flang-mod-1",
    },
    "flang-23": {
        "exe": ["flang-new-23", "flang-23"],
        "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
flags_changes = {
    "gfortran-10": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-11": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-12": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-13": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-14": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-15": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "gfortran-16": {
        "FCFLAGS": {"append": ["-fallow-invalid-boz", "-fallow-argument-mismatch"]}
    },
    "flang-7": {"FCFLAGS": {"strip": ["-g"]}},
    "flang-new-18": {
        "FCFLAGS": {
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ]
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
    },
    "flang-new-19": {
        "FCFLAGS": {
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ]
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
    },
    "flang-new-20": {
        "FCFLAGS": {
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ]
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
    },
    "flang-21": {
        "FCFLAGS": {
            "append": ["-fPIC"],  # fPIC unnecessary for static, but workaround for fpm
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ],
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
        "CC": {"replace": ["clang-21"]},
        "CXX": {"replace": ["clang++-21"]},
        "CPPFLAGS": {"strip": ["-Wdate-time"]},
    },
    "flang-22": {
        "FCFLAGS": {
            "append": ["-fPIC"],  # fPIC unnecessary for static, but workaround for fpm
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ],
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
        "CC": {"replace": ["clang-22"]},
        "CXX": {"replace": ["clang++-22"]},
        "CPPFLAGS": {"strip": ["-Wdate-time"]},
    },
    "flang-23": {
        "FCFLAGS": {
            "append": ["-fPIC"],  # fPIC unnecessary for static, but workaround for fpm
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ],
        },
        "LDFLAGS": {"strip": ["-mbranch-protection=standard"]},
        "CC": {"replace": ["clang-23"]},
        "CXX": {"replace": ["clang++-23"]},
        "CPPFLAGS": {"strip": ["-Wdate-time"]},
    },
    "lfortran": {
        "FCFLAGS": {
            "strip": [
                "-mbranch-protection=standard",
                "-fstack-protector-strong",
                "-fstack-clash-protection",
                "-ffile-prefix-map%",
                "-fcf-protection",
            ],
            "append" : [ "--cpp" ] ,
        }
    },
}


@functools.cache
def rework_flags(flavor: str, host_arch: str) -> list[str]:
    """edit env flags.
    TODO: this is the dpkg env, user may have changed
    """
    if flavor not in flags_changes:
        # cli.verbose_print(f"No env flag changes needed for {flavor}")
        return ""
    old_flags = check_output(["dpkg-buildflags"]).strip().decode("utf-8")
    old_env, new_env = {}, {}
    for line in old_flags.split("\n"):
        i = line.index("=")
        key, val = line[:i], line[i + 1 :]
        old_env[key] = val
    for var in flags_changes[flavor]:
        if "replace" in flags_changes[flavor][var]:
            new_env[var] = " ".join(flags_changes[flavor][var]["replace"])
            continue
        flags_list = old_env[var].split()
        if "strip" in flags_changes[flavor][var]:
            for r in flags_changes[flavor][var]["strip"]:
                if "%" in r:
                    prefix = r[: r.index("%")]
                    for x in flags_list:
                        if x.startswith(prefix):
                            flags_list.remove(x)
                else:
                    if r in flags_list:
                        flags_list.remove(r)
        if "append" in flags_changes[flavor][var]:
            flags_list += flags_changes[flavor][var]["append"]
        new_env[var] = " ".join(flags_list)
    return new_env


@functools.cache
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"
    if f in default_compilers:
        f = default_compilers[f]
    fullpath = which(f)
    if fullpath is None:
        fullpath = which(default_compilers['gfortran'])
    if fullpath is None:
        if orig_fc == None:
            orig_fc='gfortran' # for better error message
        raise Exception(f"fc compiler {orig_fc} is broken; bad symlink or missing install?")
    fc = Path(fullpath).resolve().name
    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}")


@functools.cache
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:
    flags = rework_flags(flavor, host_arch)
    if "LDFLAGS" in flags:
        ldflags = flags["LDFLAGS"]
    else:
        ldflags = os.environ["LDFLAGS"] if 'LDFLAGS'in os.environ else ""
    fcflags = flags["FCFLAGS"] + " " + ldflags
    fmoddir = get_fmoddir(flavor, host_arch)
    fcflags += f" -I{fmoddir}"

    return fcflags


def get_ld_flags(flavor: str, host_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
    TODO: rework in more general format
    """
    flags = rework_flags(flavor, host_arch)
    flibdir = get_flibdir(flavor, host_arch)
    if "LDFLAGS" in flags:
        ldflags = flags["LDFLAGS"]
    else:
        ldflags = os.environ["LDFLAGS"] if 'LDFLAGS'in os.environ else ""
    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?

    """
    if preferred is not None:
        return preferred
    if "PREFERRED" in os.environ:
        return os.environ["PREFERRED"]
    f = flavor if flavor in compilers else get_fc_default()
    return get_fc_default(f)


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:
        if f is None:
            f='gfortran' # for better error message
        raise Exception(f"f77 compiler {f} is broken; bad symlink or missing?")
    return Path(fullpath).resolve().name


@functools.cache
def get_fc_default(fc=None) -> str:
    """Return the default Fortran compiler"""
    if "FC_DEFAULT" in os.environ:
        return os.environ["FC_DEFAULT"]
    if fc is not None:
        abi = get_abi_vendor(fc)
        if abi in default_compilers:
            return default_compilers[abi].lower()
    default, _ = get_fc_flavor_arch(fc)
    return default


@functools.cache
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 "" in flavors:
            flavors.remove("")
        if d in flavors:
            flavors.remove(d)
        return flavors
    else:
        compilers_present = set()
        # NOTE Disable search functionality until dh-fortran is more mature ...
        #  its surprising behaviour and unwanted right now
        #
        # 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


@functools.cache
def get_fc_required(fc=None) -> set[str]:
    """Return the list of compilers that must work and libraries ship.
    Will include $(FC_DEFAULT), not overlap with $(FC_OPTIONAL)
    TODO: complete, error check
    """
    reqd = set()
    reqd.add(get_fc_default(fc))
    if "FC_REQUIRED" in os.environ:
        for f in os.environ["FC_REQUIRED"].split(" "):
            reqd.add(f)
    return reqd


@functools.cache
def is_multifortran_build():
    r = get_fc_required()
    o = get_fc_optional()
    return len(r) + len(o) > 1


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, host_arch = get_fc_flavor_arch(flavor)
    fc_default = get_fc_default(fc)
    fc_optional = " ".join(get_fc_optional(fc))
    fc_required = " ".join(get_fc_required(fc))
    fc_flags = get_fc_flags(fc, host_arch)
    fc_ldflags = get_ld_flags(fc, host_arch)
    result = f"""FC={fc}
FC_DEFAULT={fc_default}
FC_OPTIONAL={fc_optional}
FC_REQUIRED={fc_required}
PKG_CONFIG_PATH={get_pkg_config_path(fc)}
CMAKE_MODULE_PATH={get_cmake_path(fc)}
FCFLAGS={fc_flags}
FFLAGS={fc_flags}
LDFLAGS={fc_ldflags}
DESTDIR={os.getcwd()}/debian/tmp-{fc}
FPM_FC={fc}
FPM_FFLAGS={fc_flags}
FPM_LDFLAGS={fc_ldflags}"""
    new_env = rework_flags(flavor, host_arch)
    for var in new_env:
        if var in {"FCFLAGS", "LDFLAGS"}:
            continue
        result += f"\n{var}={new_env[var]}"
    return result


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",  #  (big endian) unsupported 
            "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 ? Is this even needed, or will
    # fc as flavor do?

    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, but lowercase

    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"])
