# Copyright 2022, Jerry James
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the
#    distribution.
# 3. Neither the name of Red Hat nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import string
import sys
from enum import Enum, auto
from typing import final, Tuple

# Version of this script
version=1

#
# BUILDROOT CATEGORIZATION
#

# Directories to ignore when generating %dir entries
root_dirs: set[str] = {
    '/',
    '/etc',
    '/usr',
    '/usr/bin',
    '/usr/lib',
    '/usr/lib/ocaml',
    '/usr/lib/ocaml/caml',
    '/usr/lib/ocaml/stublibs',
    '/usr/lib/ocaml/threads',
    '/usr/lib64',
    '/usr/lib64/ocaml',
    '/usr/lib64/ocaml/caml',
    '/usr/lib64/ocaml/stublibs',
    '/usr/lib64/ocaml/threads',
    '/usr/libexec',
    '/usr/sbin',
    '/usr/share',
    '/usr/share/doc'
}

def find_buildroot_toplevel(buildroot: str) -> list[str]:
    """Find toplevel files and directories in the buildroot.

    :param str buildroot: path to the buildroot
    :return: a list of toplevel files and directories in the buildroot
    """
    bfiles: list[str] = []
    for path, dirs, files in os.walk(buildroot):
        for i in range(len(dirs) - 1, -1, -1):
            d = os.path.join(path, dirs[i])[len(buildroot):]
            if d not in root_dirs and not d.startswith('/usr/share/man'):
                bfiles.append(d)
                del dirs[i]
        for f in files:
            realfile = os.path.join(path, f)[len(buildroot):]
            if realfile.startswith('/usr/share/man'):
                bfiles.append(realfile + '*')
            else:
                bfiles.append(realfile)
    return bfiles

# File suffixes that go into a devel subpackage
dev_suffixes: set[str] = {
    'a', 'cmo', 'cmt', 'cmti', 'cmx', 'cmxa', 'h', 'idl', 'ml', 'mli', 'o'
}

def is_devel_file(filename: str) -> bool:
    """Determine whether a file belongs to a devel subpackage.

    :param str filename: the filename to check
    :return: True if the file belongs to a devel subpackage, else False
    """
    return (filename == 'dune-package' or filename == 'opam' or
            (os.path.splitext(filename)[1][1:] in dev_suffixes
             and not filename.endswith('_top_init.ml')))

def find_buildroot_all(buildroot: str, devel: bool, add_star: bool) -> list[set[str]]:
    """Find all files and directories in the buildroot and optionally
    categorize them as 'main' or 'devel'.

    :param Namespace args: parsed command line arguments
    :param bool devel: True to split into 'main' and 'devel', False otherwise
    :param bool add_star: True to add a star to man page filenames
    :return: a list of files and directories, in this order: main files,
        main directories, devel files, and devel directories
    """
    bfiles: list[set[str]] = [set(), set(), set()]
    bdirs: set[str] = set()
    for path, dirs, files in os.walk(buildroot):
        for d in dirs:
            realdir = os.path.join(path, d)[len(buildroot):]
            if realdir not in root_dirs and not realdir.startswith('/usr/share/man'):
                bdirs.add(realdir)
        for f in files:
            realfile = os.path.join(path, f)[len(buildroot):]
            if devel and is_devel_file(os.path.basename(realfile)):
                bfiles[2].add(realfile)
            else:
                if add_star and realfile.startswith('/usr/share/man'):
                    bfiles[0].add(realfile + '*')
                else:
                    bfiles[0].add(realfile)
                parentdir = os.path.dirname(realfile)
                if parentdir in bdirs:
                    bfiles[1].add(parentdir)
                    bdirs.remove(parentdir)
    bfiles.append(bdirs)
    return bfiles

#
# INSTALL FILE LEXER AND PARSER
#

class TokenType(Enum):
    """The types of tokens that can appear in an opam *.install file."""
    ERROR  = auto()
    EOF    = auto()
    COLON  = auto()
    LBRACE = auto()
    RBRACE = auto()
    LBRACK = auto()
    RBRACK = auto()
    STRING = auto()
    FIELD  = auto()

@final
class InstallFileLexer:
    """Convert an opam *.install file into a sequence of tokens."""
    __slots__ = ['index', 'text']

    def __init__(self) -> None:
        """Create an empty opam *.install file lexer."""
        self.text = ''
        self.index = 0

    def lex_file(self, filename: str) -> None:
        """Prepare to read tokens from an opam *.install file.

        :param str filename: the name of the file to read from
        """
        with open(filename, 'r') as f:
            # Limit reads to 4 MB in case this file is bogus.
            # Most install files are under 4K.
            self.text = f.read(4194304)
        self.index = 0

    def skip_whitespace(self) -> None:
        """Skip over whitespace in the input."""
        while self.index < len(self.text) and \
              (self.text[self.index] == '#' or
               self.text[self.index] in string.whitespace):
            if self.text[self.index] == '#':
                while (self.index < len(self.text) and
                       self.text[self.index] != '\n' and
                       self.text[self.index] != '\r'):
                    self.index += 1
            else:
                self.index += 1

    def token(self) -> Tuple[TokenType, str]:
        """Get the next token from the opam *.install file.

        :return: a pair containing the type and text of the next token
        """
        self.skip_whitespace()
        if self.index < len(self.text):
            ch = self.text[self.index]
            if ch == ':':
                self.index += 1
                return (TokenType.COLON, ch)
            if ch == '{':
                self.index += 1
                return (TokenType.LBRACE, ch)
            if ch == '}':
                self.index += 1
                return (TokenType.RBRACE, ch)
            if ch == '[':
                self.index += 1
                return (TokenType.LBRACK, ch)
            if ch == ']':
                self.index += 1
                return (TokenType.RBRACK, ch)
            if ch == '"':
                start = self.index + 1
                end = start
                while end < len(self.text) and self.text[end] != '"':
                    end += 2 if self.text[end] == '\\' else 1
                self.index = end + 1
                return (TokenType.STRING, self.text[start:end])
            if ch in string.ascii_letters:
                start = self.index
                end = start + 1
                while (end < len(self.text) and
                       (self.text[end] == '_' or
                        self.text[end] in string.ascii_letters)):
                    end += 1
                self.index = end
                return (TokenType.FIELD, self.text[start:end])
            return (TokenType.ERROR, ch)
        else:
            return (TokenType.EOF, '')

@final
class InstallFileParser:
    """Parse opam *.install files to generate RPM %files lists."""

    __slots__ = ['buildroot', 'lexer', 'libdir', 'pkgs']

    def __init__(self, buildroot: str, libdir: str, devel: bool) -> None:
        """Initialize an OCaml .install file parser.

        :param str buildroot: path to the buildroot
        :param str libdir: the OCaml library directory
        :param bool devel: True to split into main and devel packages
        """
        self.buildroot = find_buildroot_all(buildroot, devel, False)
        self.libdir = libdir
        self.lexer = InstallFileLexer()
        self.pkgs: dict[str, set[str]] = dict()

    def add_package(self, pkgname: str, filename: str) -> None:
        """Add a mapping from pkgname to filename.

        :param str pkgname: the package that acts as the map key
        :param str filename: the filename to add to the package set
        """
        if pkgname not in self.pkgs:
            self.pkgs[pkgname] = set()
        self.pkgs[pkgname].add(filename)

    def register_file(self, pkgname: str, filename: str) -> None:
        """Register one file listed in an opam *.install file.

        :param str pkgname: name of the package to which this file belongs
        :param str filename: name of the file
        """
        if filename in self.buildroot[0]:
            if filename.startswith('/usr/share/man'):
                self.add_package(pkgname, filename + '*')
            else:
                self.add_package(pkgname, filename)
                dirname = os.path.dirname(filename)
                if dirname in self.buildroot[1]:
                    self.add_package(pkgname, '%dir ' + dirname)
                    self.buildroot[1].remove(dirname)
        elif filename in self.buildroot[2]:
            if filename.startswith('/usr/share/man'):
                self.add_package(pkgname + '-devel', filename + '*')
            else:
                self.add_package(pkgname + '-devel', filename)
                dirname = os.path.dirname(filename)
                if dirname in self.buildroot[3]:
                    self.add_package(pkgname + '-devel', '%dir ' + dirname)
                    self.buildroot[3].remove(dirname)

    def parse_file(self, filename: str) -> None:
        """Parse a .install file and add the contents to internal file lists.
        If there are any parse errors, we assume this file is not really an
        opam .install file and abandon the parse.

        :param str filename: name of the .install file to parse
        """
        # Get the package name from the filename
        pkgname = os.path.splitext(os.path.basename(filename))[0]

        # Map opam installer names to directories
        opammap: dict[str, str] = {
            'lib': os.path.join(self.libdir, pkgname),
            'lib_root': self.libdir,
            'libexec': os.path.join(self.libdir, pkgname),
            'libexec_root': self.libdir,
            'bin': '/usr/bin',
            'sbin': '/usr/sbin',
            'toplevel': os.path.join(self.libdir, 'toplevel'),
            'share': os.path.join('/usr/share', pkgname),
            'share_root': '/usr/share',
            'etc': os.path.join('/etc', pkgname),
            'doc': os.path.join('/usr/doc', pkgname),
            'stublibs': os.path.join(self.libdir, 'stublibs'),
            'man': '/usr/share/man'
        }

        # Prepare the lexer
        self.lexer.lex_file(filename)

        # Parse the file
        toktyp, token = self.lexer.token()
        while toktyp == TokenType.FIELD:
            libname = token
            toktyp, token = self.lexer.token()
            if toktyp != TokenType.COLON:
                return

            toktyp, token = self.lexer.token()
            if toktyp != TokenType.LBRACK:
                return

            directory = opammap.get(libname)
            if not directory:
                return

            toktyp, token = self.lexer.token()
            while toktyp == TokenType.STRING:
                nexttp, nexttk = self.lexer.token()
                if nexttp == TokenType.LBRACE:
                    nexttp, nexttk = self.lexer.token()
                    if nexttp == TokenType.STRING:
                        filnam = os.path.join(directory, nexttk)
                        bracetp, bractk = self.lexer.token()
                        if bracetp != TokenType.RBRACE:
                            return
                        nexttp, nexttk = self.lexer.token()
                    else:
                        return
                elif libname == 'man':
                    index = token.rfind('.')
                    if index < 0:
                        return
                    mandir = os.path.join(directory, 'man' + token[index+1:])
                    filnam = os.path.join(mandir, os.path.basename(token))
                else:
                    filnam = os.path.join(directory, os.path.basename(token))
                toktyp, token = nexttp, nexttk
                self.register_file(pkgname, filnam)

            if toktyp != TokenType.RBRACK:
                return
            toktyp, token = self.lexer.token()

        if toktyp != TokenType.EOF:
            return

    def parse_files(self) -> dict[str, set[str]]:
        """Find all .install files and parse each one.

        For some projects, there are install files in both the project root
        directory and somewhere under "_build", so be careful not to parse the
        same install file twice.
        :return: a map from package names to set of files to install
        """
        install_files = set()
        for path, dirs, files in os.walk('.'):
            for f in files:
                if f.endswith('.install') and f not in install_files:
                    install_files.add(f)
                    self.parse_file(os.path.join(path, f))
        return self.pkgs

#
# MAIN INTERFACE
#

def ocaml_files(no_devel: bool, separate: bool, buildroot: str, libdir: str) -> None:
    """Generate %files lists from an installed buildroot.

    :param bool no_devel: False to split files into a main package and a devel
        package
    :param bool separate: True to place each OCaml module in an RPM package
    :param str buildroot: the installed buildroot
    :param str libdir: the OCaml library directory
    """
    if separate:
        parser = InstallFileParser(buildroot, libdir, not no_devel)
        pkgmap = parser.parse_files()
        for pkg in pkgmap:
            with open('.ofiles-' + pkg, 'w') as f:
                for entry in pkgmap[pkg]:
                    f.write(entry + '\n')
    elif no_devel:
        with open('.ofiles', 'w') as f:
            for entry in find_buildroot_toplevel(buildroot):
                f.write(entry + '\n')
    else:
        files = find_buildroot_all(buildroot, True, True)
        with open('.ofiles', 'w') as f:
            for entry in files[0]:
                f.write(entry + '\n')
            for entry in files[1]:
                f.write('%dir ' + entry + '\n')
        with open('.ofiles-devel', 'w') as f:
            for entry in files[2]:
                f.write(entry + '\n')
            for entry in files[3]:
                f.write('%dir ' + entry + '\n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Support for building OCaml RPM packages')
    parser.add_argument('-n', '--no-devel',
                        action='store_true',
                        default=False,
                        help='suppress creation of a devel subpackage')
    parser.add_argument('-s', '--separate',
                        action='store_true',
                        default=False,
                        help='separate packaging.  Each OCaml module is in a distinct RPM package.  All modules are in a single RPM package by default.')
    parser.add_argument('-v', '--version',
                        action='version',
                        version=f'%(prog)s {str(version)}')
    parser.add_argument('buildroot', help='RPM build root')
    parser.add_argument('libdir', help='OCaml library directory')
    args = parser.parse_args()
    ocaml_files(args.no_devel, args.separate, args.buildroot, args.libdir)
