# Copyright 2022 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Manages the list of Pigweed modules.

Used by modules.gni to generate:

- a build arg for each module,
- a list of module paths (pw_modules),
- a list of module tests (pw_module_tests), and
- a list of module docs (pw_module_docs).
"""

import argparse
import difflib
import enum
import io
import os
from pathlib import Path
import sys
import subprocess
from typing import Iterator, Optional, Sequence

_COPYRIGHT_NOTICE = '''\
# Copyright 2022 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.'''

_WARNING = '\033[31m\033[1mWARNING:\033[0m '  # Red WARNING: prefix
_ERROR = '\033[41m\033[37m\033[1mERROR:\033[0m '  # Red background ERROR: prefix

_MISSING_MODULES_WARNING = (
    _WARNING
    + '''\
The PIGWEED_MODULES list is missing the following modules:
{modules}

If the listed modules are Pigweed modules, add them to PIGWEED_MODULES.

If the listed modules are not actual Pigweed modules, remove any stray pw_*
directories in the Pigweed repository (git clean -fd).
'''
)

_OUT_OF_DATE_WARNING = (
    _ERROR
    + '''\
The generated Pigweed modules list .gni file is out of date!

Regenerate the modules lists and commit it to fix this:

  ninja -C {out_dir} update_modules

  git add {file}
'''
)

_FORMAT_FAILED_WARNING = (
    _ERROR
    + '''\
Failed to generate a valid .gni from PIGWEED_MODULES!

This may be a Pigweed bug; please report this to the Pigweed team.
'''
)

_DO_NOT_SET = 'DO NOT SET THIS BUILD ARGUMENT!'


def _module_list_warnings(root: Path, modules: Sequence[str]) -> Iterator[str]:
    missing = _missing_modules(root, modules)
    if missing:
        yield _MISSING_MODULES_WARNING.format(
            modules=''.join(f'\n  - {module}' for module in missing)
        )

    if any(modules[i] > modules[i + 1] for i in range(len(modules) - 1)):
        yield _WARNING + 'The PIGWEED_MODULES list is not sorted!'
        yield ''
        yield 'Apply the following diff to fix the order:'
        yield ''
        yield from difflib.unified_diff(
            modules,
            sorted(modules),
            lineterm='',
            n=1,
            fromfile='PIGWEED_MODULES',
            tofile='PIGWEED_MODULES',
        )

        yield ''


def _generate_modules_gni(
    prefix: Path, modules: Sequence[str]
) -> Iterator[str]:
    """Generates a .gni file with variables and lists for Pigweed modules."""
    script_path = Path(__file__).resolve()
    script = script_path.relative_to(script_path.parent.parent).as_posix()

    yield _COPYRIGHT_NOTICE
    yield ''
    yield '# Build args and lists for all modules in Pigweed.'
    yield '#'
    yield f'# DO NOT EDIT! Generated by {script}.'
    yield '#'
    yield '# To add modules here, list them in PIGWEED_MODULES and build the'
    yield '# update_modules target and commit the updated version of this file:'
    yield '#'
    yield '#   ninja -C out update_modules'
    yield '#'
    yield '# DO NOT IMPORT THIS FILE DIRECTLY!'
    yield '#'
    yield '# Import it through //build_overrides/pigweed.gni instead.'
    yield ''
    yield '# Declare a build arg for each module.'
    yield 'declare_args() {'

    for module in modules:
        module_path = prefix.joinpath(module).as_posix()
        yield f'dir_{module} = get_path_info("{module_path}", "abspath")'

    yield '}'
    yield ''
    yield '# Declare these as GN args in case this is imported in args.gni.'
    yield '# Use a separate block so variables in the prior block can be used.'
    yield 'declare_args() {'
    yield f'# A list with paths to all Pigweed module. {_DO_NOT_SET}'
    yield 'pw_modules = ['

    for module in modules:
        yield f'dir_{module},'

    yield ']'
    yield ''

    yield f'# A list with all Pigweed module test groups. {_DO_NOT_SET}'
    yield 'pw_module_tests = ['

    for module in modules:
        yield f'"$dir_{module}:tests",'

    yield ']'
    yield ''
    yield f'# A list with all Pigweed modules docs groups. {_DO_NOT_SET}'
    yield 'pw_module_docs = ['

    for module in modules:
        yield f'"$dir_{module}:docs",'

    yield ']'
    yield ''
    yield '}'


def _missing_modules(root: Path, modules: Sequence[str]) -> Sequence[str]:
    return sorted(
        frozenset(
            str(p.relative_to(root)) for p in root.glob('pw_*') if p.is_dir()
        )
        - frozenset(modules)
    )


class Mode(enum.Enum):
    WARN = 0  # Warn if anything is out of date
    CHECK = 1  # Fail if anything is out of date
    UPDATE = 2  # Update the generated modules lists


def _parse_args() -> dict:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument('root', type=Path, help='Root build dir')
    parser.add_argument('modules_list', type=Path, help='Input modules list')
    parser.add_argument('modules_gni_file', type=Path, help='Output .gni file')
    parser.add_argument(
        '--mode', type=Mode.__getitem__, choices=Mode, required=True
    )
    parser.add_argument(
        '--stamp',
        type=Path,
        help='Stamp file for operations that should only run once (warn)',
    )
    return vars(parser.parse_args())


def main(
    root: Path,
    modules_list: Path,
    modules_gni_file: Path,
    mode: Mode,
    stamp: Optional[Path] = None,
) -> int:
    """Manages the list of Pigweed modules."""
    prefix = Path(os.path.relpath(root, modules_gni_file.parent))
    modules = modules_list.read_text().splitlines()

    # Detect any problems with the modules list.
    warnings = list(_module_list_warnings(root, modules))
    errors = []

    modules.sort()  # Sort in case the modules list in case it wasn't sorted.

    # Check if the contents of the .gni file are out of date.
    if mode in (Mode.WARN, Mode.CHECK):
        text = io.StringIO()
        for line in _generate_modules_gni(prefix, modules):
            print(line, file=text)

        process = subprocess.run(
            ['gn', 'format', '--stdin'],
            input=text.getvalue().encode('utf-8'),
            stdout=subprocess.PIPE,
        )
        if process.returncode != 0:
            errors.append(_FORMAT_FAILED_WARNING)

        # Make a diff of required changes
        modules_gni_relpath = os.path.relpath(modules_gni_file, root)
        diff = list(
            difflib.unified_diff(
                modules_gni_file.read_text().splitlines(),
                process.stdout.decode('utf-8', errors='replace').splitlines(),
                fromfile=os.path.join('a', modules_gni_relpath),
                tofile=os.path.join('b', modules_gni_relpath),
                lineterm='',
                n=1,
            )
        )
        # If any differences were found, print the error and the diff.
        if diff:
            errors.append(
                _OUT_OF_DATE_WARNING.format(
                    out_dir=os.path.relpath(os.curdir, root),
                    file=os.path.relpath(modules_gni_file, root),
                )
            )
            errors.append('Expected Diff:\n')
            errors.append('\n'.join(diff))
            errors.append('\n')

    elif mode is Mode.UPDATE:  # Update the modules .gni file
        with modules_gni_file.open('w', encoding='utf-8') as file:
            for line in _generate_modules_gni(prefix, modules):
                print(line, file=file)

        process = subprocess.run(
            ['gn', 'format', modules_gni_file], stdout=subprocess.DEVNULL
        )
        if process.returncode != 0:
            errors.append(_FORMAT_FAILED_WARNING)

    # If there are errors, display them and abort.
    if warnings or errors:
        for line in warnings + errors:
            print(line, file=sys.stderr)

        # Delete the stamp so this always reruns. Deleting is necessary since
        # some of the checks do not depend on input files.
        if stamp and stamp.exists():
            stamp.unlink()

        if mode is Mode.WARN:
            return 0

        if mode is Mode.CHECK:
            return 1

        return 1 if errors else 0  # Allow warnings but not errors when updating

    if stamp:
        stamp.touch()

    return 0


if __name__ == '__main__':
    sys.exit(main(**_parse_args()))
