#!/usr/bin/python3
"""
Module file handling for Fortran

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

import click
import magic
import os
import dhfortran.debhelper as dh
import dhfortran.cli as cli
import dhfortran.compilers as cmplrs
import gzip
from subprocess import check_output
import fnmatch

# cache
flang_modfile_hashes = {}


def guess_flang_version(fc, modfile):
    """ flang now ships with "v1" as a format version, but hardcodes a dep on modules
    via a hash. This essentially makes different compiler versions incompatible.
    Try figuring version from module hashes
    """

    def get_hashes_from_modfile(filename , version = None):
        """ Parse a flang v1 modfile, getting hashes,names """        
        with open(filename,'rt') as f:
            lines = f.readlines()
        hash, hashes = '', {}
        for line in lines:
            match line:
                case l if l.find('!mod$ v1 sum:') != -1:
                    hash = l[l.rfind(':')+1:-1]
                case l if l.startswith('!need$'):
                    _,hash,_,name = l.split(' ')
                    hashes[hash] = (version, name[:-1])
                case l if l.startswith('module'):
                    _,name = l.split(' ')
                    hashes[hash] = (version, name[:-1])
                    return hashes, hash
                case _:
                    # something strange happened
                    return hashes, hash

    if fc is None and 'FC' in os.environ:
        fc = os.environ['FC']
    if  fc is not None:
        if fc.startswith('flang-'):
            mod_version = fc[6:] # got the version
            return mod_version
                
    # Do it the hard way.
    if flang_modfile_hashes == {}:
        # Build db from system modules
        flang_versions = set([ key[-2:] for key in cmplrs.compilers.keys()
                         if key.startswith('flang-')])
        for fv in flang_versions:
            for root,dir,files in os.walk(f"/usr/lib/llvm-{fv}/include/flang"):
                for m in fnmatch.filter(files,"*.mod"):
                    hashes, _ = get_hashes_from_modfile(f"{root}/{m}",fv)
                    flang_modfile_hashes.update(hashes)
                    
    hashes, myhash = get_hashes_from_modfile(modfile)
    if myhash in flang_modfile_hashes:
        myver,_ = flang_modfile_hashes[myhash]
        for h in hashes:
            flang_modfile_hashes[h] = (myver, hashes[h][1])
        return myver
    # failed, just use default
    return cmplrs.default_compilers['flang'][6:]
            

    
def which_compiler_flavor_and_mod_version(fc: str, modfile: str) -> tuple[str, str]:
    """For a modfile (path), return a tuple of strings for the compiler
    and modversion"""

    # Helper .
    def get_first_n_lines(modfile,n=2 ):
        try:
            magick = magic.from_file(modfile)
        except Exception as ex:
            raise Exception(f"Can't open modfile {modfile}: {ex}")
        if magick.startswith("gzip compressed"):
            with gzip.open(modfile, "rt") as f:
                return [f.readline(), f.readline()]
        if "text" in magick:
            with open(modfile, "rt") as f:
                # Strip any Unicode BOMs
                return [f.readline().replace("\ufeff", ""), f.readline()]
        lines = check_output(["strings", modfile]).decode("utf-8").split("\n")
        return lines[0:n]

 
    lines = get_first_n_lines(modfile,2)
    match lines[0]:
        case line if line.startswith("GFORTRAN"):
            mod_version = line.split()[3][1:-1]
            if line.find("created from flang") != -1:
                compiler = "flangext"
            else:
                compiler = "gfortran"
        case line if line.startswith("LCompilers Modfile"):
            compiler, mod_version = "lfortran", "0" # lines[1]
        case line if line.startswith("V35"):
            compiler, mod_version = "nvhpc", "35"
        case line if line.startswith("V"):
            # could also be PGI, but we don't ship PGI on Debian ... its an old flang
            compiler = "flang"
            mod_version = line.split()[1][1:]
        case line if line.startswith("!mod$"):
            compiler = "flang"
            mod_version = guess_flang_version(fc, modfile)
        case line if line.startswith("k"):
            # Intel: k820309, 2021.13.0
            # TODO check if version-specific
            compiler, mod_version = "intel", "0"
        case _:
            raise Exception(f"Can't parse file: {modfile}")

    return (compiler, mod_version)


##
## Modfile-specific interface for debhelper
##
class ModFileHelper(dh.DhFortranHelper):

    def __init__(self, options: dict):
        super().__init__(options, "fortran-mod", "dh_fortran_mod")
        # For tracking which compilers we need
        self.modversions = {}

    def can_work_automatically(self, *cmdline_files) -> bool:
        """Should we automatically automagically detect mod files, 
        or work on specified packages?

        Run if we haven't run already, haven't been passed cmdline files,
        """
        if len(cmdline_files) > 0:
            return False
        for p in self.packages:
            if self.has_acted_on_package(p):
                return False
        if len(self.config_file_pkgs) != 0:
            cli.verbose_print(
                "dh_fortran won't run automatically: has specified fortran-mod files"
            )
            return False
        if len(self.libdevel_pkgs) != 1:
            return False
        pkg = self.libdevel_pkgs[0]
        cli.verbose_print(
            f"No fortran-mod files found,1 devel package {pkg}:"
            "may install mod files automatically"
        )
        return True

    def autogenerate_pkg_filelist(
        self, seachdirs: list[str], dest=None
    ) -> list[tuple[str, str, str]]:
        """Return list of (pkg, file_pattern) tuple to work on.
        We walk the search path looking for files to operate on
        """
        mfiles = []
        for d in seachdirs:
            if d is None:
                continue
            for root, dir, files in os.walk(d):
                for item in fnmatch.filter(files, "*.mod"):
                    mfiles.append(f"{root}/{item}")
                for item in fnmatch.filter(files, "*.smod"):
                    mfiles.append(f"{root}/{item}")

        return [(self.libdevel_pkgs[0], f, dest) for f in mfiles]

    def compute_dest(self, modfile, target_pkg, target_dest=None):
        """Where does modfile go ?"""
        # cli.debug_print("compute_dest [mod]  {modfile} {target_dest}")
        # Should be called by base DebHelper() to move files
        cmplr, version = which_compiler_flavor_and_mod_version(self.options.fc, modfile)
        fmoddir = f"/usr/lib/{cmplrs.multiarch}/fortran/{cmplr}-mod-{version}"
        if target_dest is None:
            return fmoddir
        if target_dest.startswith("/"):
            return target_dest
        else:
            return f"{fmoddir}/{target_dest}"

    def process_file(self, pkg, oldfile, target_pkg, target_dest=None):
        """Called by DebHelper
        Module file gets copied; metadata gets stored"""
        cli.debug_print(
            f"process_file [module]  {oldfile=} {target_pkg=} {target_dest=}"
        )
        try:
            compiler, version = which_compiler_flavor_and_mod_version(self.options.fc,oldfile)
        except Exception as ex:
            cli.verbose_print(f"skipping unknown file {oldfile}")
            return
        if pkg not in self.modversions:
            self.modversions[pkg] = set()
        self.modversions[pkg].add(f"{compiler}-mod-{version}")
        dest = self.compute_dest(oldfile, target_pkg, target_dest)
        assert target_pkg != '' and target_pkg is not None
        dh.install_dir(f"debian/{target_pkg}/{dest}")
        os.chmod(oldfile, 0o644)
        self.doit(["cp", oldfile, f"debian/{target_pkg}/{dest}"])
        #if oldfile.startswith(self.options.tmpdir):
        #    self.doit(["rm", oldfile])


@click.command(
    context_settings=dict(
        ignore_unknown_options=False,
    )
)
@click.option("--fc", help="Fortran compiler flavor (eg gfortran-14, flang-21)")
@click.argument("files", nargs=-1, type=click.UNPROCESSED)
@cli.debhelper_common_args
def dh_fortran_mod(files, *args, **kwargs):
    """
    B<dh_fortran_mod> is a debhelper program that finds Fortran module 
    and submodule files and adds dependencies to B<gfortran-$version> 
    as required to the package using via the variables
    B<$fortran:Depends} and B<${fortran:Recommends}>.

    B<dh_fortran_mod> is expected to be automatically added using 
    the debhelper "addon" B<fortran_mod> ie. either automatically, 
    by build-depending on 'dh-sequence-fortran', or explicitly:

        dh $@ --with fortran


    B<dh_fortran_mod>Searches the debian/ directory for files 
    B<debian/pkg.fortran-mod>$ which list module files to include, 
    with the same syntax as debhelper install files.

    =head1 OPTIONS

    =over 4

    =item B<--tmpdir=>I<dir>

    Look in the specified directory for files to be installed (by default debian/tmp)

    Typically Fortran module files are included in library development packages.

    =back

    =head1 TODO

    Add dh-fortran-mod support for generic fortran compilers (ifx, etc).

    =over 4

    B<dh_fortran_mod> will be expanded to find mod files automatically from the 
    I<debian/tmp> directory.
    It will enable the installation of mod files in parallel for multiple compilers.
    It will install .smod files for Fortran 2018.

    The fortran-mod file syntax follows dh_install: pairs of sources and optional
    target directories.  The default directory will be $fmoddir 
    ( /usr/lib/$multiarch/fortran/$compiler_mod_directory/)
    If the target directory is absolute (starts with a  '/'), this directory is used
    in the target package.  If the target  does not absolute, it will be treated 
    as a subdirectory of $fmoddir.

    Currently four flavours of Fortran compiler are supported: gfortran-*  ('gfortran'),
    flang-new-* ('flang') , flang-to-external-fc-* ('flangext') and 
    lfortran ('lfortran').

    $compiler_mod_directory is based on the compiler module version: 
    currently gfortran-mod-16 for gfortran-15 (and older), flang-mod-1
    for flang-new-18+ and lfortran-mod-0 for lfortran.

    For flang-to-external-fc-* the version is flangext-mod-15 (assuming 
    gfortran-15 as the external compiler); in principle flang-to-external-fc 
    ('flangext' flavour) and gfortran are intercompatible but intermixing is avoided.

    These will be updated for incompatible compiler versions.
        
    """
    cli.debug_print(f"dh_fortran_mod called with files {files} kwargs {kwargs}")

    d = ModFileHelper(dh.build_options(**kwargs))

    searchdirs = [d.options.tmpdir, "."]
    d.process_and_move_files(searchdirs, *files)
    depends = []
    recommends = []
    fc_default = cmplrs.get_fc_default()
    fc_optional = cmplrs.get_fc_optional()
    for package in d.packages:
        if package in d.modversions:
            for m in sorted(d.modversions[package]):
                default_comps = set()
                optional_comps = set()
                for c in cmplrs.compilers:
                    if cmplrs.compilers[c]["mod"] == m:
                        if c  ==   fc_default:
                            default_comps.add(c)
                        if c in fc_optional:
                            optional_comps.add(c)
                if r:= "| ".join(sorted(default_comps)):
                    depends.append(f"{r} | {m}")
                if r:= "| ".join(sorted(optional_comps)):
                    recommends.append(f"{r} | {m}")
            d.addsubstvar(package, "fortran:Depends", ",".join(depends))
            d.addsubstvar(package, "fortran:Recommends", ",".join(recommends))
        if "config_file" in d.packages[package] and package in d.modversions:
            # config file defined, so create cleanup script
            for modversion in d.modversions[package]:
                d.autoscript(
                    package,
                    "postrm",
                    "postrm-fortran-mod",
                    {
                        "MULTIARCH": cmplrs.multiarch,
                        "FCOMPILERMOD": modversion,
                    },
                )
    d.save_substvars()


if __name__ == "__main__":
    import pytest

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